문제 설명
(u, v) 형태의 간선 목록이 주어지며, 이 간선들은 하나의 트리를 구성한다고 가정해 보겠습니다. 이때 각 간선마다 해당 간선을 지나는 고유한 경로의 총 개수를 구하고, 결과를 입력된 순서 그대로 반환해야 합니다.
예를 들어 edges = [[0, 1], [0, 2], [1, 3], [1, 4]]가 입력으로 주어진다면,
출력은 [6, 4, 4, 4]가 됩니다.
알고리즘 접근 방법
이 문제를 해결하기 위해 다음 단계를 따릅니다 −
주어진 간선들로부터 인접 리스트(adj)를 생성합니다.
count := 빈 맵(딕셔너리)을 준비합니다.
x와 parent를 인자로 받는 dfs() 함수를 정의합니다.
count[x] := 1로 초기화합니다.
adj[x]에 있는 각 이웃 nb에 대해 다음을 수행합니다.
nb가 parent와 같다면 건너뜁니다.
count[x] := count[x] + dfs(nb, x)
count[x]를 반환합니다.
메인 로직에서는 다음을 수행합니다 −
dfs(0, -1)을 호출합니다.
ans := 새로운 리스트를 생성합니다.
edges의 각 간선 (a, b)에 대해 다음을 수행합니다.
x := count[a]와 count[b] 중 최솟값
ans의 끝에 (x * (count[0] - x))를 추가합니다.
ans를 반환합니다.
동작 원리
핵심 아이디어는 간단합니다. 트리에서 임의의 간선 하나를 제거하면 트리는 정확히 두 개의 연결 요소로 분리됩니다. 한쪽 서브트리에 속한 노드 수를 x, 전체 노드 수를 n이라 하면 반대쪽에는 n - x개의 노드가 남습니다. 어떤 경로가 해당 간선을 지나려면 양쪽 컴포넌트에서 노드를 하나씩 골라 연결해야 하므로, 그 간선을 포함하는 고유한 경로의 수는 x × (n - x)가 됩니다.
DFS를 한 번 수행하며 각 노드의 서브트리 크기를 미리 계산해 두면, 모든 간선에 대해 상수 시간(O(1)) 안에 답을 구할 수 있습니다. 따라서 전체 시간 복잡도는 O(n)으로 매우 효율적입니다.
예제 코드
더 나은 이해를 위해 다음 구현을 살펴보겠습니다 −
from collections import defaultdict class Solution: def solve(self, edges): adj = defaultdict(list) for a, b in edges: adj[a].append(b) adj[b].append(a) count = defaultdict(int) def dfs(x, parent): count[x] = 1 for nb in adj[x]: if nb == parent: continue count[x] += dfs(nb, x) return count[x] dfs(0, -1) ans = [] for a, b in edges: x = min(count[a], count[b]) ans.append(x * (count[0] - x)) return ans ob = Solution() edges = [ [0, 1], [0, 2], [1, 3], [1, 4] ] print(ob.solve(edges))
입력
[ [0, 1], [0, 2], [1, 3], [1, 4] ]
출력
[6, 4, 4, 4]