무방향이면서 연결된 하나의 트리가 있고, 이 트리에는 N개의 노드가 존재한다고 가정해 봅시다. 노드는 0부터 N-1까지 라벨이 붙어 있으며, N-1개의 간선이 주어집니다. i번째 간선은 edges[i][0]과 edges[i][1] 두 노드를 연결합니다. 우리가 구해야 하는 것은 ans[i]가 '노드 i와 나머지 모든 노드 사이의 거리의 합'을 나타내는 리스트입니다.
예를 들어 입력이 N = 6이고 edges = [(0,1),(0,2),(2,3),(2,4),(2,5)]라면, 출력은 [8,12,6,10,10,10]이 됩니다.
문제 접근 방식
모든 노드 쌍 사이의 거리를 일일이 계산하면 O(N²) 시간이 걸려 비효율적입니다. 대신 DFS를 두 번 수행하는 재배치(re-rooting) 기법을 사용하면 O(N) 시간에 문제를 해결할 수 있습니다.
핵심 아이디어는 다음과 같습니다.
- 1차 DFS(dfs1): 임의의 루트(0번 노드)에서 시작하여, 각 서브트리에 포함된 노드의 개수(cnt)와 해당 서브트리 내부의 거리 합(ans)을 자식 노드부터 계산해 올라옵니다.
- 2차 DFS(dfs2): 부모 노드의 답을 이용해 자식 노드의 답을 도출합니다. 루트에서 자식으로 이동할 때, 자식 쪽 서브트리에 속한 cnt[child]개의 노드는 1씩 가까워지고, 나머지 N - cnt[child]개의 노드는 1씩 멀어집니다.
이제 단계별로 살펴보겠습니다.
dfs1() 함수 정의
dfs1(node, parent)는 다음과 같이 동작합니다.
- graph[node]의 모든 인접 노드를 순회합니다.
- 인접 노드(child)가 parent와 같지 않다면:
- dfs1(child, node)를 재귀 호출합니다.
- cnt[node] += cnt[child] — 자식 서브트리의 노드 수를 누적합니다.
- ans[node] += cnt[child] + ans[child] — 자식 서브트리 내부 거리 합과, node에서 자식 서브트리 각 노드까지의 거리(각각 1씩 증가)를 더합니다.
dfs2() 함수 정의
dfs2(node, parent)는 다음과 같이 동작합니다.
- graph[node]의 모든 인접 노드를 순회합니다.
- 인접 노드(child)가 parent와 같지 않다면:
- ans[child] = ans[node] - cnt[child] + (N - cnt[child]) — 부모의 값을 기준으로 자식의 거리 합을 재계산합니다.
- dfs2(child, node)를 재귀 호출합니다.
메인 로직
- 멤버 변수로 정수 배열 ans, cnt와 크기 10005의 인접 리스트 graph를 선언합니다.
- N을 저장하고, ans는 크기 N으로, cnt는 모든 값을 1로 초기화하여 크기 N 배열로 생성합니다.
- 간선 정보를 순회하며 graph[u]에 v를, graph[v]에 u를 추가해 양방향 그래프를 구성합니다.
- dfs1(0, -1)과 dfs2(0, -1)을 차례로 호출한 뒤 ans를 반환합니다.
아래 예시 코드를 통해 더 잘 이해해 보겠습니다.
예제 코드 (C++)
#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
cout << "[";
for(int i = 0; i<v.size(); i++){
cout << v[i] << ", ";
}
cout << "]"<<endl;
}
class Solution {
public:
void dfs1(int node, int parent) {
for (int i = 0; i < graph[node].size(); i++) {
int child = graph[node][i];
if (child != parent) {
dfs1(child, node);
cnt[node] += cnt[child];
ans[node] += cnt[child] + ans[child];
}
}
}
void dfs2(int node, int parent) {
for (int i = 0; i < graph[node].size(); i++) {
int child = graph[node][i];
if (child != parent) {
ans[child] = ans[node] - cnt[child] + N - cnt[child];
dfs2(child, node);
}
}
}
vector<int> ans;
vector<int> cnt;
vector<int> graph[10005];
int N;
vector<int> sumOfDistancesInTree(int N, vector<vector<int> >& edges) {
this->N = N;
ans = vector<int>(N);
cnt = vector<int>(N, 1);
int n = edges.size();
for (int i = 0; i < n; i++) {
int u = edges[i][0];
int v = edges[i][1];
graph[u].push_back(v);
graph[v].push_back(u);
}
dfs1(0, -1);
dfs2(0, -1);
return ans;
}
};
main(){
Solution ob;
vector<vector<int>> v = {{0,1},{0,2},{2,3},{2,4},{2,5}};
print_vector(ob.sumOfDistancesInTree(6, v));
}입력
{{0,1},{0,2},{2,3},{2,4},{2,5}}출력
[8, 12, 6, 10, 10, 10]
복잡도 분석
- 시간 복잡도: O(N) — 각 노드를 두 번의 DFS에서 한 번씩만 방문합니다.
- 공간 복잡도: O(N) — 인접 리스트, ans, cnt 배열 및 재귀 호출 스택이 필요합니다.
마무리
이 문제는 트리에서 루트를 바꿔가며(re-rooting) 거리 합을 효율적으로 갱신하는 대표적인 기법을 보여줍니다. 첫 번째 DFS로 한 루트 기준의 답을 구하고, 두 번째 DFS에서 부모-자식 관계를 활용해 나머지 노드들의 답을 상수 시간에 유도하는 것이 핵심입니다. 이 패턴은 트리 관련 다양한 최적화 문제에 응용될 수 있으므로 꼭 익혀두시길 권장합니다.