Pythonでn分木の最長パスの長さを求めるプログラムの書き方

各要素が (u, v) という形式を持ち、u が v の親であることを表す辺リストが与えられているとします。このとき、木の中で最も長いパスの長さを求める必要があります。ここでいうパスの長さとは、「そのパスに含まれるノードの総数 + 1」のことです。

たとえば、入力が下図のような n 分木だった場合を考えてみましょう。

この場合の出力は 5 になります。なぜなら、パス [1, 4, 5, 7] には合計 4 つのノードが含まれており、パスの長さは 1 + 4 = 5 となるからです。

解き方のアプローチ

この問題は、幅優先探索(BFS)を2回実行するという定番テクニックで効率よく解けます。まず任意のノードから BFS を行って最も遠いノード f を見つけ、次に f を起点として再度 BFS を実行します。2 回目の BFS で得られる最大距離が、まさに木の直径(最長パスの長さ)になります。

具体的な手順は以下の通りです。

• 与えられた辺リストから、グラフの隣接リスト g を作成します
• 各ノードまでの距離を記録するためのマップ d を用意します
• BFS を行う関数 bfs(o) を定義します(引数 o は開始ノード)
  • d[o] := 1 として、開始ノードの距離を 1 に設定します
  • f := o として、最遠ノードの候補を初期化します
  • q := [o] として、探索用キューを初期化します
  • キュー内の各ノード x について、隣接ノード y を順に調べます
    • y が未訪問であれば、d[y] := d[x] + 1 とします
    • d[y] > d[f] であれば、f := y と更新します
    • y をキューに追加します
  • 最後に、最も遠いノード f を返します
• メイン処理では、任意のノード o から bfs(o) を呼び出して最遠ノード f を取得し、マップ d をリセットした上で bfs(f) を再び実行します。その結果得られるノードの距離 d[bfs(f)] が答えとなります
• グラフが空の場合は 0 を返します

Pythonでの実装例

それでは、実際のコードを見てみましょう。

def solve(edges):
    g = {}
    for u, v in edges:
        if u not in g:
            g[u] = []
        g[u] += (v,)
        if v not in g:
            g[v] = []
        g[v] += (u,)
    d = {}

    def bfs(o):
        d[o] = 1
        f = o
        q = [o]
        for x in q:
            for y in g[x]:
                if y not in d:
                    d[y] = d[x] + 1
                    if d[y] > d[f]:
                        f = y
                    q += (y,)
        return f

    for o in g:
        f = bfs(o)
        d = {}
        return d[bfs(f)]
    return 0

edges = [(1, 2),(1, 3),(1, 4),(4, 5),(5,7),(1,6),(4,8)]
print(solve(edges))

入力

[(1, 2),(1, 3),(1, 4),(4, 5),(5,7),(1,6),(4,8)]

出力

5

アルゴリズムのポイント

なぜ BFS を 2 回行うだけで最長パスが求まるのでしょうか。木において「任意のノードから最も遠いノード」は、必ず直径(最長パス)の端点のどちらかになっていることが証明されています。したがって、1 回目の BFS で見つけた最遠ノード f は直径の片方の端点であり、f からもう一度 BFS を行ったときの最大距離がそのまま直径になります。

この手法の計算量は O(N)(N はノード数)です。すべてのノードペア間の経路を総当たりする O(N²) のアプローチと比べて大幅に効率的なので、ノード数が多い木でも高速に動作します。