Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 모든 아이템 구매 최소 비용 구하기 – 다익스트라 알고리즘 활용


문제 개요

0부터 N-1까지 번호가 붙은 N개의 아이템이 있다고 가정해 보겠습니다. 크기가 S인 2차원 리스트 sets가 주어지며, i번째 세트는 sets[i][2]의 가격에 구매할 수 있고, 구매 시 sets[i][0]부터 sets[i][1] 범위에 속한 모든 아이템을 한꺼번에 얻게 됩니다. 또한 크기가 N인 리스트 removals가 주어져서, removals[i]만큼의 비용을 지불하면 i번째 아이템 하나를 버릴 수 있습니다.

목표는 0부터 N-1까지의 모든 아이템을 정확히 하나씩 확보하는 최소 비용을 구하는 것이며, 만약 불가능하다면 -1을 반환해야 합니다.

예시 입력

sets = [
    [0, 4, 4],
    [0, 5, 12],
    [2, 6, 9],
    [4, 8, 10]
]
removals = [2, 5, 4, 6, 8]

이 경우 출력은 4입니다. 첫 번째 세트 [0, 4, 4]를 가격 4에 구매하면 0번부터 4번까지, 즉 모든 아이템을 한 번에 얻을 수 있기 때문입니다.

접근 방법: 그래프 최단 경로(다익스트라)

이 문제는 그래프로 모델링하면 다익스트라(Dijkstra) 최단 경로 알고리즘으로 깔끔하게 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • 각 세트 [s, e, w]는 “노드 s에서 노드 e+1로 이동하는 비용 w의 간선”으로 표현합니다. 즉, 세트를 구매하면 해당 범위의 끝 다음 위치로 점프하게 됩니다.

  • 각 제거 비용 removals[i]는 “노드 i+1에서 노드 i로 이동하는 비용 r의 간선”으로 표현합니다.

  • 노드 0에서 출발하여 노드 N에 도달하는 최단 거리가 곧 정답이 됩니다.

알고리즘 단계

  • N := removals의 길이

  • graph := 크기 (N + 1)의 인접 리스트 생성

  • sets의 각 (s, e, w)에 대해 → graph[s]에 [e+1, w] 추가

  • removals의 각 인덱스 i와 값 r에 대해 → graph[i+1]에 [i, r] 추가

  • pq := 시작점 [0, 0]을 담은 최소 힙 생성

  • dist := 기본값이 무한대인 딕셔너리, dist[0] := 0으로 초기화

  • pq가 빌 때까지 반복:

    • d, e := 힙에서 가장 작은 원소 꺼내기

    • dist[e] < d이면 이미 더 나은 경로가 존재하므로 건너뛰기

    • e == N이면 d를 반환 (목표 지점 도달)

    • graph[e]의 각 이웃 (nei, w)에 대해:

      • d2 := d + w

      • d2 < dist[nei]이면 dist[nei] := d2로 갱신하고 [d2, nei]를 힙에 push

  • 반복이 끝날 때까지 목표에 도달하지 못하면 -1 반환

구현 예제

아래 파이썬 구현을 통해 더 잘 이해할 수 있습니다.

import heapq
from collections import defaultdict

class Solution:
    def solve(self, sets, removals):
        N = len(removals)
        graph = [[] for _ in range(N + 1)]
        for s, e, w in sets:
            graph[s].append([e + 1, w])
        for i, r in enumerate(removals):
            graph[i + 1].append([i, r])
        pq = [[0, 0]]
        dist = defaultdict(lambda: float("inf"))
        dist[0] = 0
        while pq:
            d, e = heapq.heappop(pq)
            if dist[e] < d:
                continue
            if e == N:
                return d
            for nei, w in graph[e]:
                d2 = d + w
                if d2 < dist[nei]:
                    dist[nei] = d2
                    heapq.heappush(pq, [d2, nei])
        return -1

ob = Solution()
print(ob.solve([
    [0, 4, 4],
    [0, 5, 12],
    [2, 6, 9],
    [4, 8, 10]
], [2, 5, 4, 6, 8]))

입력

[[0, 4, 4],
[0, 5, 12],
[2, 6, 9],
[4, 8, 10]], [2, 5, 4, 6, 8]

출력

4

복잡도 분석

노드 수를 V = N + 1, 간선 수를 E = S + N이라 할 때, 다익스트라 알고리즘의 시간 복잡도는 우선순위 큐 연산이 지배하므로 O(E log V)입니다. 공간 복잡도는 그래프 저장과 거리 맵에 필요한 O(V + E)입니다.