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

Python으로 가중 그래프에서 최소 비용 경로 찾기

문제 개요

2차원 정수 리스트 edges가 하나의 무방향 그래프(undirected graph)를 나타낸다고 가정해 보겠습니다. 입력의 각 행은 [u, v, w] 형태의 간선 정보를 담고 있으며, 이는 노드 u와 v가 서로 연결되어 있고 해당 간선의 가중치가 w임을 의미합니다. 그래프는 0부터 n-1까지 총 n개의 노드로 구성됩니다.

여기서 경로의 비용은 다음과 같이 정의됩니다.

경로 비용 = (경로에 포함된 간선의 개수) × (경로상 간선 가중치의 최댓값)

우리가 구해야 할 것은 노드 0에서 출발하여 노드 n-1에 도달하는 경로 중 최소 비용입니다. 만약 두 노드 사이에 어떤 경로도 존재하지 않는다면 -1을 반환해야 합니다.

예시

입력이 다음과 같다고 해보겠습니다.

edges = [
    [0, 2, 100],
    [1, 2, 200],
    [1, 3, 100],
    [2, 3, 300]
]

이 경우 출력은 600이 됩니다.

접근 방법

이 문제는 BFS(너비 우선 탐색)와 가중치 상한을 단계적으로 낮추는 전략을 조합하여 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • 특정 가중치 상한(weight_cap)을 정하면, 그 상한 이하의 간선만 사용해서 목적지까지 도달할 수 있는지 BFS로 확인할 수 있습니다.
  • 도달이 가능하다면, 해당 경로의 비용(간선 수 × 최대 가중치)을 계산하여 결과를 갱신합니다.
  • 이후 가중치 상한을 낮춰가며 더 작은 비용의 경로가 있는지 반복적으로 탐색합니다.

알고리즘 단계

해결 과정을 단계별로 살펴보면 다음과 같습니다.

  • 그래프 인접 리스트(graph), 간선 가중치 맵(weights), 최대 가중치(max_weight), 노드 개수(N)를 초기화합니다.
  • 각 간선 [u, v, w]에 대해 양방향 연결 정보와 가중치를 저장하고, N과 max_weight를 갱신합니다.
  • result를 무한대로 초기화한 뒤, max_weight가 0 이상인 동안 다음을 반복합니다.
    • bfs(0, max_weight)를 호출하여 도달 거리 d와 실제 사용된 최대 가중치 weight를 얻습니다.
    • d가 0 이상이면 result를 min(result, d × weight)로 갱신하고, max_weight를 weight - 1로 낮춥니다.
    • 도달하지 못하면 반복을 종료합니다.
  • result가 무한대가 아니면 result를, 그렇지 않으면 -1을 반환합니다.

BFS 함수의 동작

  • 목표 노드(target)는 N - 1입니다.
  • 큐 Q에는 (노드, 거리, 현재까지의 최대 가중치) 튜플을 저장합니다.
  • 방문 여부 배열 visited로 중복 탐색을 방지합니다.
  • 큐에서 요소를 꺼냈을 때 목표 노드에 도달했다면 거리와 최대 가중치를 반환합니다.
  • 인접 노드 중 아직 방문하지 않았고, 간선 가중치가 weight_cap 이하인 노드만 큐에 추가합니다. 이때 누적 최대 가중치를 함께 갱신합니다.
  • 모든 탐색이 끝나도 목표에 도달하지 못하면 (-1, -1)을 반환합니다.

구현 코드

아래는 위 알고리즘을 Python으로 구현한 전체 코드입니다.

from collections import defaultdict, deque

class Solution:
    def solve(self, edges):
        graph = defaultdict(list)
        weights = {}
        max_weight = 0
        N = 0
        for u, v, w in edges:
            graph[u].append(v)
            graph[v].append(u)
            weights[u, v] = w
            weights[v, u] = w
            N = max(N, u + 1, v + 1)
            max_weight = max(max_weight, w)

        def bfs(root, weight_cap):
            target = N - 1
            Q = deque([(root, 0, 0)])
            visited = [False] * N
            visited[0] = True
            while Q:
                v, d, current_weight = Q.pop()
                if v == N - 1:
                    return d, current_weight
                for w in graph[v]:
                    if visited[w]:
                        continue
                    new_weight = weights[v, w]
                    if new_weight <= weight_cap:
                        visited[w] = True
                        Q.appendleft((w, d + 1, max(current_weight, new_weight)))
            return -1, -1

        result = float("inf")
        while max_weight >= 0:
            d, weight = bfs(0, max_weight)
            if d >= 0:
                result = min(result, d * weight)
                max_weight = weight - 1
            else:
                break
        return result if result < float("inf") else -1

ob = Solution()
print(ob.solve([
    [0, 2, 100],
    [1, 2, 200],
    [1, 3, 100],
    [2, 3, 300]
]))

실행 결과

입력

[
    [0, 2, 100],
    [1, 2, 200],
    [1, 3, 100],
    [2, 3, 300]
]

출력

600

동작 원리 분석

예제에서 처음 max_weight는 300입니다. 이 상태에서 BFS를 수행하면 노드 0 → 2 → 3 경로를 통해 목적지에 도달할 수 있으며, 간선 2개와 최대 가중치 300을 사용하므로 비용은 2 × 300 = 600입니다. 이후 max_weight를 299로 낮춰 다시 탐색하면, 가중치 300짜리 간선을 사용할 수 없어 더 이상 목적지에 도달할 수 없습니다. 따라서 최종 결과는 600이 됩니다.

이처럼 가중치 상한을 점진적으로 줄여가며 BFS를 반복하는 방식은, 경로 비용의 정의상 "간선 개수"와 "최대 가중치"라는 두 요소의 곱으로 결정되기 때문에 효과적으로 동작합니다. 시간 복잡도는 O(E² × V) 수준으로, 간선 가중치 종류만큼 BFS를 반복하게 됩니다.