C++で木の直径を求めるアルゴリズムを解説

木の直径とは

無向木(undirected tree)が与えられたとき、その直径を求めることを考えます。木の直径とは、木の中で最も長い経路に含まれる辺の数のことです。

ここでは、木は辺のリストとして与えられます。edges[i] = [u, v] は、ノードuとノードvをつなぐ双方向の辺を表します。また、各ノードには {0, 1, ..., edges.length} の集合からラベルが割り当てられています。

例として、次のような木を考えてみましょう。

この場合、最も長い経路の長さは4となるため、出力は4になります。

解法のアプローチ

木の直径を効率的に求めるには、DFS(深さ優先探索)を2回実行するという有名な手法を使います。具体的な手順は以下の通りです。

アルゴリズムの手順

1. マップ l を定義します。
2. dfs() というメソッドを定義します。引数としてノード v、訪問済みフラグ配列 visited、グラフ graph、現在の深さ c を受け取ります。
3. visited[v] を true に設定し、ans を 0 で初期化します。
4. graph[v] の隣接ノードをすべて確認し、未訪問のノードがあれば再帰的に dfs() を呼び出し、その結果の最大値を ans に記録します。
5. 現在の深さ c が best より大きければ、best を c に、node を v に更新します(最も遠いノードを記録)。
6. visited[v] を false に戻します(バックトラック)。
7. c と ans の最大値を返します。
8. メインの処理では、辺リスト e を受け取り、n を e のサイズとし、サイズ n+1 の隣接リスト graph を作成します。
9. すべての辺について、双方向に隣接関係を登録します。
10. サイズ n+1 の訪問済み配列 visited と visited2 を作成し、best と node を 0 で初期化します。
11. まずノード0から dfs(0, visited, graph) を呼び出し、最も遠いノード node を特定します。
12. そのノードから 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

アルゴリズムのポイント

この手法が正しく機能する理由は、任意のノードから最も遠いノードは、必ず直径を構成する経路の端点のいずれかになるという木の性質に基づいています。

そのため、1回目のDFSで端点の候補となるノードを見つけ、2回目のDFSでそのノードから最遠ノードまでの距離を測ることで、木全体の直径を正確に求められます。計算量はノード数をNとすると O(N) で、非常に効率的なアルゴリズムです。