Python으로 무방향 가중치 그래프에서 주어진 비용 이하의 경로가 존재하는지 확인하는 방법

문제 이해하기

가중치가 부여된 무방향 그래프(undirected graph)가 주어졌다고 가정해 보겠습니다. 우리는 두 개의 정점과 비용 상한값 limit을 입력으로 받아, 해당 비용 이하의 비용으로 두 정점을 연결하는 경로가 존재하는지 확인하는 query() 함수를 구현해야 합니다. 경로가 존재하면 True를, 존재하지 않으면 False를 반환합니다.

예제

다음과 같은 그래프가 있다고 가정합니다.

쿼리가 (0, 2, 10), (3, 1, 30), (4, 3, 30)일 때 출력 결과는 다음과 같습니다.

False
True
True

첫 번째 쿼리 → False: 정점 0에서 정점 2로 가는 비용 10 이하의 경로가 존재하지 않습니다.
두 번째 쿼리 → True: 정점 3에서 정점 1로 가는 비용 10짜리 경로가 존재하며, 이는 30보다 작습니다.
세 번째 쿼리 → True: 정점 4에서 정점 3으로 가는 비용 30짜리 경로가 존재하며, 이는 30과 같습니다.

접근 방법: 유니온-파인드(Union-Find) 활용

이 문제는 간선의 가중치를 기준으로 그래프의 연결 상태가 어떻게 변화하는지 추적하는 방식으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

모든 간선을 가중치 기준으로 오름차순 정렬합니다.
간선을 하나씩 추가하면서 유니온-파인드로 정점들을 묶어 가고, 서로 다른 가중치의 간선이 모두 반영된 시점마다 현재 연결 상태(각 정점의 루트)를 스냅샷으로 저장합니다.
쿼리가 들어오면 이분 탐색(bisect)으로 limit 미만의 가중치만 사용했을 때의 스냅샷을 찾아, 두 정점이 같은 집합에 속하는지 확인합니다.

알고리즘 단계

weights: 그래프에 존재하는 서로 다른 가중치 값들을 담은 리스트
connections: 각 가중치 시점별 정점 연결(루트) 정보를 담은 리스트
query(p, q, limit) 함수 정의:
- index := 정렬된 순서를 유지하면서 limit이 삽입될 수 있는 weights 내 위치 (bisect_left 사용)
- index가 0이라면, 즉 limit보다 작은 가중치의 간선이 하나도 없다면 False를 반환합니다.
- connections[index-1]에서 정점 p와 q의 루트가 같으면 True, 아니면 False를 반환합니다.

구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

import bisect
class Solution(object):

   def __init__(self, n, edgeList):
      def find(node):
         if parent[node]!=node:
            parent[node] = find(parent[node])
         return parent[node]

      def union(x,y):
         parent[find(y)] = find(x)
         return

      parent = {i:i for i in range(n)}
      edgeList.sort(key = lambda x:x[2])
      self.connections = []
      self.weights = []
      for index,(i,j,weight) in enumerate(edgeList):
         union(i,j)
         if index!=len(edgeList)-1 and weight == edgeList[index+1][2]:
            continue
         self.weights.append(weight)
         self.connections.append([find(i) for i in parent])


   def query(self, p, q, limit):
      index = bisect.bisect_left(self.weights,limit)
      if index==0:
         return False
      return self.connections[index-1][p] == self.connections[index-1][q]

ob = Solution(5, [[0, 1, 10], [0, 2, 20], [1, 4, 10], [0, 3, 10], [1, 2, 20], [2, 3, 10]])
print(ob.query(0, 2, 10))
print(ob.query(3, 1, 30))
print(ob.query(4, 3, 30))

입력

ob = Solution(5, [[0, 1, 10], [0, 2, 20], [1, 4, 10], [0, 3, 10], [1, 2, 20], [2, 3, 10]])
print(ob.query(0, 2, 10))
print(ob.query(3, 1, 30))
print(ob.query(4, 3, 30))

출력

False
True
True

코드 동작 원리

초기화(__init__): 모든 정점을 자기 자신을 루트로 하는 독립적인 집합으로 초기화한 뒤, 간선 목록을 가중치순으로 정렬합니다. 간선을 순서대로 유니온하며, 동일한 가중치를 가진 간선들이 모두 처리된 시점에 현재 연결 상태(각 정점의 루트)를 connections에 저장하고 해당 가중치를 weights에 기록합니다.

쿼리 처리(query): bisect_left를 사용해 limit 미만의 가중치만으로 만들어진 마지막 스냅샷을 찾습니다. 두 정점의 루트가 동일하다는 것은 해당 비용 범위 내에서 두 정점이 이미 연결되어 있음을 의미합니다.

이 방식은 전처리에 O(E log E), 각 쿼리 처리에 O(log W)(W는 서로 다른 가중치의 수)가 소요되므로, 동일한 그래프에 대해 수많은 쿼리를 처리해야 하는 상황에서 특히 효율적입니다.