【Python】木の直径

木の直径

木の単純パスの長さの最大値の事を、木の直径と言います。
例えば、以下のグラフであれば、

頂点１と４（または５）を繋ぐ単純パスが最も長いものであるので、このグラフの直径は３です。

$O(N^2)$の解法

頂点間のパスが最も長いものを見つければ良いので、BFSやDFS等を用いて、全ての頂点間の距離を求めてあげれば、その最大値が木の直径になります。ただし、これでは$N \sim 10^5$程度になると解けません。

$O(N)$の解法

木の直径を求めることは、$O(N)$で可能です。最短経路を2回求める事で、これが達成できます。

• 頂点１からの最短経路が最も長い頂点$u$を求める。
• 頂点$u$からの最短経路が最も長い頂点$v$を求める。
• 頂点$u, v$間の距離が木の直径となる。

最短経路を求めるアルゴリズムは、ダイクストラ法などがありますが、今回は木であるため、わざわざダイクストラ法を使う必要はありません。（閉路が無いため）単純にDFSやBFSで、端の頂点まで見ていくだけで、$O(N)$で最短経路が求められます。

例題

Atcoder 競プロ典型 90 問 第3問

与えられるグラフは木です。木であれば、任意の隣接していない頂点間を繋ぐ事で、必ずその辺を含む閉路が作れます。この問題は、木の直径を求めて、その頂点間を繋ぐ辺を追加することで、スコアを最大化することができます。
今回は、DFSを用いて解答します。

import sys
sys.setrecursionlimit(10**6)

N = int(input())

edges = [[] for _ in range(N)]
for _ in range(N-1):
    a, b = map(int, input().split())
    edges[a-1].append(b-1)
    edges[b-1].append(a-1)

dist = [float('inf')] * N
dist[0] = 0
max_dist = {'node': -1, 'dist': -1}

def dfs(x, prev):
    for y in edges[x]:
        if y == prev:
            continue
        dist[y] = dist[x] + 1
        if max_dist['dist'] < dist[y]:
            max_dist['dist'] = dist[y]
            max_dist['node'] = y
        dfs(y, x)

dfs(0, -1)
dist = [float('inf')] * N
node = max_dist['node']
dist[node] = 0
max_dist = {'node': -1, 'dist': -1}
dfs(node, -1)

print(max_dist['dist'] + 1)

木には閉路が存在しないため、探索中に同じ頂点を通ることはありません。従って、まだ探索していない頂点を見つければ、その頂点に移動して貪欲に最短距離を更新していけば良いです。

参考

https://twitter.com/e869120/status/1377752658149175299