C++로 트리의 지름(Diameter) 구하기 — DFS 두 번으로 최장 경로 찾기

문제 개요

무방향 트리가 하나 주어졌을 때, 이 트리의 지름(diameter)을 구하는 것이 목표입니다. 트리의 지름이란 트리 안에서 가장 긴 경로에 포함된 간선의 개수를 의미합니다.

트리는 간선 리스트 형태로 주어지며, edges[i] = [u, v]는 노드 u와 노드 v를 잇는 양방향 간선을 나타냅니다. 각 노드의 레이블은 {0, 1, ..., edges.length} 범위에 속합니다.

예를 들어 다음과 같은 트리가 있다고 가정해 보겠습니다.

이 경우 가장 긴 경로는 3 → 2 → 1 → 4 → 5이며, 여기에는 간선이 4개 포함되어 있으므로 출력값은 4가 됩니다.

접근 방법: DFS 두 번 수행하기

이 문제는 깊이 우선 탐색(DFS)을 두 번 수행하는 방식으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

트리의 임의의 노드에서 출발해 DFS로 가장 멀리 있는 노드를 찾으면, 그 노드는 반드시 지름 경로의 양 끝점 중 하나입니다. 따라서 첫 번째 DFS로 끝점 후보를 찾고, 그 노드에서 두 번째 DFS를 수행하면 실제 지름을 구할 수 있습니다.

구체적인 알고리즘 단계는 다음과 같습니다.

1. 맵 l을 정의합니다.
2. dfs() 메서드를 정의합니다. 이 메서드는 노드 v, 방문 여부 배열 visited, 그래프, 현재 깊이 c를 인자로 받으며 다음과 같이 동작합니다.
   • visited[v] := true로 설정하고 ans := 0으로 초기화합니다.
   • 0부터 graph[v]의 크기까지 반복하면서, 아직 방문하지 않은 인접 노드라면 dfs(graph[v][i], visited, graph, c + 1)를 재귀 호출하고 그 결과로 ans를 갱신합니다.
   • c > best라면 best := c, node := v로 갱신합니다.
   • visited[v] := false로 되돌려 백트래킹합니다.
   • max(c, ans)를 반환합니다.
3. 메인 메서드에서는 간선 리스트 e를 입력받습니다.
4. n := e의 크기로 설정하고, 크기가 n + 1인 graph 배열을 생성합니다.
5. 0부터 n − 1까지 반복하면서 graph[e[i][0]]에 e[i][1]을, graph[e[i][1]]에 e[i][0]을 삽입해 양방향 간선을 구성합니다.
6. 크기 n + 1의 visited 배열과 visited2 배열을 만들고, best := 0, node := 0으로 초기화합니다.
7. dfs(0, visited, graph)를 호출해 노드 0에서 가장 멀리 있는 노드를 찾아냅니다.
8. dfs(node, visited2, graph)를 호출한 결과를 반환합니다. 이 값이 곧 트리의 지름입니다.

C++ 구현 예시

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

#include <bits/stdc++.h>
using namespace std;
#define pb push_back
class Solution {
public:
    map <int ,int > l;
    int best;
    int node;
    int dfs(int v, bool* visited, vector <int> graph[], int c = 0){
       visited[v] = true;
       int ans = 0;
       for(int i = 0; i < graph[v].size(); i++){
          if(!visited[graph[v][i]])ans = max(ans,dfs(graph[v][i], visited, graph, c+1));
       }
       if(c > best){
          best = c;
          node = v ;
       }
       visited[v] = false;
       return max(c,ans);
    }
    int treeDiameter(vector<vector<int>>& e) {
       int n = e.size();
       vector <int> graph[n+1];
       for(int i = 0; i < n; i++){
          graph[e[i][0]].pb(e[i][1]);
          graph[e[i][1]].pb(e[i][0]);
       }
       bool* visited = new bool[n+1]();
       best = 0;
       node = 0;
       dfs(0, visited, graph);
       bool* visited2 = new bool[n+1]();
       return dfs(node, visited2, graph);
    }
};
main(){
    vector<vector<int>> v = {{0,1},{1,2},{2,3},{1,4},{4,5}};
    Solution ob;
    cout <<ob.treeDiameter(v);
}

입력

[[0,1],[1,2],[2,3],[1,4],[4,5]]

출력

4

결과 분석

주어진 트리에서 첫 번째 DFS는 노드 0에서 시작해 가장 멀리 있는 노드(노드 3)를 찾아냅니다. 이후 노드 3에서 두 번째 DFS를 수행하면 노드 5까지의 경로(3 → 2 → 1 → 4 → 5)가 가장 길다는 것을 알 수 있고, 이 경로의 간선 개수는 4개이므로 최종 결과는 4가 됩니다.

이 방법은 전체 트리를 두 번만 순회하면 되므로 시간 복잡도는 O(N)이며, 노드 개수 N에 대해 매우 효율적으로 동작합니다.