문제 개요
n개의 노드로 이루어진 무방향 가중치 그래프가 edgeList로 주어집니다. edgeList[i]는 세 개의 값 (u, v, w)을 가지며, 이는 u에서 v로 연결되는 거리 w의 경로가 존재함을 의미합니다.
또한 별도의 queries 배열이 주어지는데, queries[i]는 (p, q, lim) 형태입니다. 이 질의는 p에서 q까지 직접 또는 다른 노드를 거쳐 도달할 수 있는 경로 중, 총 거리가 lim보다 작은 경로가 존재하는지를 묻습니다. 우리는 각 질의에 대해 True/False 결과를 담은 배열을 반환해야 합니다.
입력 예시
예를 들어 다음과 같은 그래프가 입력으로 주어졌다고 가정해 보겠습니다.
이 경우 출력은 [True, False, True]가 됩니다. 그 이유는 다음과 같습니다.
- 첫 번째 질의: 1에서 4로 가려면 1 → 3 → 4 경로를 따라 비용 11로 이동할 수 있으므로 True입니다.
- 두 번째 질의: 2에서 3으로 3 미만의 비용으로 이동할 수 없으므로 False입니다.
- 세 번째 질의: 1 → 3 → 2 경로를 통해 비용 14로 이동할 수 있고, 이는 15보다 작으므로 True입니다.
풀이 접근법
이 문제는 유니온-파인드(Union-Find) 알고리즘을 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 간선을 가중치 기준으로 오름차순 정렬하고, 질의 역시 거리 제한(lim) 기준으로 정렬한 뒤, 제한값보다 작은 가중치의 간선들을 순서대로 하나의 집합으로 합쳐 나가는 것입니다. 두 노드가 같은 집합에 속해 있다면, 해당 제한 내에서 두 노드를 잇는 경로가 존재한다는 의미입니다.
구체적인 단계는 다음과 같습니다.
- parent 배열을 0부터 n까지 초기화합니다.
- rank 배열을 크기 n+1로 만들고 0으로 채웁니다.
- find(parent, x) 함수를 정의합니다. x의 루트 부모를 찾아 반환하며, 경로 압축(path compression)을 적용합니다.
- union(parent, a, b) 함수를 정의합니다. a와 b의 루트를 찾은 뒤, rank를 비교하여 두 트리를 병합합니다.
- 메인 로직에서는 다음을 수행합니다.
- edgeList를 가중치(w) 기준으로 정렬합니다.
- res 배열을 질의 개수만큼 생성하고 0으로 초기화합니다.
- queries에 원래 인덱스 i를 붙인 후, 거리 제한(lim) 기준으로 정렬합니다.
- 포인터 ind를 0으로 초기화합니다.
- 각 질의 (a, b, w)에 대해, edgeList[ind]의 가중치가 w보다 작은 동안 해당 간선을 union으로 병합하고 ind를 증가시킵니다.
- find(parent, a)와 find(parent, b)가 같은지 비교한 결과를 res[i]에 저장합니다.
- res 배열을 반환합니다.
구현 코드
다음 구현을 통해 더 자세히 이해해 보겠습니다.
def solve(n, edgeList, queries):
parent = [i for i in range(n+1)]
rank = [0 for i in range(n+1)]
def find(parent, x):
if parent[x] == x:
return x
parent[x] = find(parent, parent[x])
return parent[x]
def union(parent, a, b):
a = find(parent, a)
b = find(parent, b)
if a == b:
return
if rank[a] < rank[b]:
parent[a] = b
elif rank[a] > rank[b]:
parent[b] = a
else:
parent[b] = a
rank[a] += 1
edgeList.sort(key=lambda x: x[2])
res = [0] * len(queries)
queries = [[i, ch] for i, ch in enumerate(queries)]
queries.sort(key=lambda x: x[1][2])
ind = 0
for i, (a, b, w) in queries:
while ind < len(edgeList) and edgeList[ind][2] < w:
union(parent, edgeList[ind][0], edgeList[ind][1])
ind += 1
res[i] = find(parent, a) == find(parent, b)
return res
n = 4
edgeList = [(1,2,16),(1,3,8),(2,4,3),(2,3,6),(4,3,3)]
queries = [(1,4,12),(2,3,3),(1,2,15)]
print(solve(n, edgeList, queries))입력
4, [(1,2,16),(1,3,8),(2,4,3),(2,3,6),(4,3,3)],[(1,4,12),(2,3,3),(1,2,15)]
출력
[True, False, True]
시간 복잡도 분석
간선 정렬에는 O(E log E), 질의 정렬에는 O(Q log Q)의 시간이 소요됩니다. 이후 유니온-파인드 연산은 경로 압축과 rank 기반 병합 덕분에 사실상 상수 시간에 처리되므로, 전체 시간 복잡도는 O((E + Q) log(E + Q))로 매우 효율적입니다. 이는 각 질의마다 BFS나 DFS로 경로를 탐색하는 방식(O(Q × (V + E)))보다 훨씬 빠릅니다.