문제 소개
n개의 정점과 m개의 간선으로 구성된 가중치 그래프가 있다고 가정해 봅시다. 모든 간선의 가중치는 2의 거듭제곱(1, 2, 4, 8 등) 형태로 주어집니다. 그래프는 완전히 연결되어 있어 어떤 정점에서든 다른 정점으로 이동할 수 있으며, 두 정점 간의 이동 비용은 경로에 포함된 모든 간선 가중치의 합입니다. 이때 우리가 구해야 하는 것은 모든 정점 쌍 사이의 최소 비용의 총합입니다.
예를 들어 아래와 같은 그래프가 입력으로 주어지고,
정점의 개수 n = 6이라면 출력 결과는 2696이 됩니다. 즉, 모든 정점 쌍 간의 최단 거리를 모두 더한 값이 2696인 것입니다.
핵심 아이디어
간선 가중치가 모두 2의 거듭제곱일 때는 크루스칼(Kruskal) 알고리즘으로 만든 최소 신장 트리(MST)가 모든 정점 쌍 간의 최단 경로를 보장합니다. 따라서 먼저 MST를 구성한 뒤, 트리의 각 간선이 몇 개의 정점 쌍 경로에 사용되는지 세면 됩니다. 특정 간선을 제거했을 때 트리가 child개와 n - child개의 정점을 가진 두 부분으로 나뉜다면, 해당 간선은 child × (n − child)개의 경로에 포함되며, 각 경로마다 간선 가중치 2^w만큼의 비용이 추가됩니다.
해결 절차
이 문제는 다음 단계를 통해 해결할 수 있습니다.
- par_finder() 함수 정의 — 인자로 i와 par 배열을 받습니다.
- par[i]가 -1이면 i를 반환합니다(자기 자신이 루트라는 의미).
- res := par_finder(par[i], par)로 재귀적으로 루트를 찾습니다.
- 경로 압축을 위해 par[i] := res로 갱신합니다.
- res를 반환합니다.
- helper() 함수 정의 — 인자로 i, par, w, G, n을 받습니다.
- child := 1로 초기화합니다.
- G[i]의 각 항목에 대해:
- item[0]이 부모(par)와 같으면 다음 반복으로 넘어갑니다.
- 그렇지 않으면 child := child + helper(item[0], i, item[1], G, n)으로 자식 서브트리의 크기를 누적합니다.
- par가 -1이 아니면(현재 노드가 루트가 아니면), ans := ans + child * (n - child) * (2^w)를 누적합니다.
- child를 반환합니다.
- G := n + 1개의 빈 리스트를 담은 새 리스트를 생성합니다.
- edges := 새 리스트를 생성하고, roads의 각 항목 (u, v, w)에 대해 (u-1, v-1, w)를 추가합니다(0 기반 인덱스로 변환).
- edges를 간선 가중치를 기준으로 오름차순 정렬합니다.
- par := 크기가 n + 1이고 모두 -1로 초기화된 새 리스트를 생성합니다.
- r_edge := 새 리스트를 생성합니다.
- edges의 각 간선 i에 대해:
- par_finder(i[0], par) == par_finder(i[1], par)이면(사이클이 발생하므로) 건너뜁니다.
- 그렇지 않으면 r_edge에 i를 추가하고, G[i[0]]에 (i[1], i[2])를, G[i[1]]에 (i[0], i[2])를 추가한 뒤 두 집합을 병합합니다: par[par_finder(i[0], par)] := par_finder(i[1], par).
- ans := 0으로 초기화합니다.
- helper(0, -1, 0, G, n)을 호출합니다.
- ans를 반환합니다.
구현 예시
아래 구현 예시를 통해 더 자세히 이해해 보겠습니다.
def par_finder(i, par) : if par[i] == -1 : return i res = par_finder(par[i], par) par[i] = res return res def helper(i, par, w, G, n) : global ans child = 1 for item in G[i] : if item[0] == par : continue else : child += helper(item[0],i,item[1], G, n) if par != -1 : ans += child * (n - child) * (1 << w) return child def solve(n, roads): global ans G = [[] for i in range(n + 1)] edges = [] for item in roads : u,v,w = map(int, item) edges.append((u-1, v-1, w)) edges = sorted(edges,key = lambda item : item[2]) par = [-1 for i in range(n + 1)] r_edge = [] for i in edges : if par_finder(i[0], par) == par_finder(i[1], par) : continue else : r_edge.append(i) G[i[0]].append((i[1],i[2])) G[i[1]].append((i[0],i[2])) par[par_finder(i[0], par)] = par_finder(i[1], par) ans = 0 helper(0, -1, 0, G, n) return ans print(solve(6, [(1,4,8), (2,4,4), (3,4,4), (3,4,2), (5,3,8), (6,3,2)]))
입력
6, [(1,4,8), (2,4,4), (3,4,4), (3,4,2), (5,3,8), (6,3,2)]
출력
2696
복잡도 분석
크루스칼 알고리즘의 간선 정렬 단계가 전체 수행 시간을 지배하므로 시간 복잡도는 O(m log m)이며, 이후 MST 위에서 진행하는 DFS 탐색은 O(n)입니다. 따라서 정점과 간선이 많은 그래프에서도 효율적으로 동작합니다.