木の直径

木の直径、重み付き木の直径のライブラリ作ってないと思って作ってみた。
python です。
操作は、幅優先探索を2回するだけ。
AtCoder Grand Contest で直径を使う問題が出たけど、
直径まで考えが及ばなかったので、覚えるために記事を書いてる。

アルゴリズム

適当な頂点から最も遠い頂点 をとる。
頂点 から最も遠い頂点 までの距離がこの木の「直径」になっている。

木の直径を求めるプログラム

def diameter(N, E):#diamete of a tree
    from collections import deque
    V = [-1]*N
    V[0] = 0
    q = deque([0])
    while q:
        v = q.popleft()
        for w in E[v]:
            if V[w] == -1:
                V[w] = V[v] + 1
                q.append(w)
    s = V.index( max(V))
    p = deque([s])
    W = [-1]*N
    W[s] = 0
    while p:
        v = p.popleft()
        for w in E[v]:
            if W[w] == -1:
                W[w] = W[v] + 1
                p.append(w)
    return max(W)

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

同じように重み付き木の直径を次のように求めることができる。

重み付き木の直径を求めるプログラム

def diameter(N, E, D):#diamete of a tree
from collections import deque
V = [-1]*N
V[0] = 0
q = deque([0])
while q:
v = q.popleft()
for w in E[v]:
if V[w] == -1:
V[w] = V[v] + D[(v,w)]
q.append(w)
s = V.index( max(V))
p = deque([s])
W = [-1]*N
W[s] = 0
while p:
v = p.popleft()
for w in E[v]:
if W[w] == -1:
W[w] = W[v] + D[(v,w)]
p.append(w)
return max(W)

from collections import defaultdict
N = int( input())
E = [ [] for _ in range(N)]
D = defaultdict( int)
for _ in range(N-1):
a, b, w = map( int, input().split())
E[a].append(b)
E[b].append(a)
D[(a,b)] = D[(b,a)] = w
print( diameter( N, E, D))

おまけ

list でも、deque と同様の操作ができるから、list でやることが多かったけど、
どうやら list は かかるけど、deque は でできるらしい。
公式ドキュメントにあった。
collections --- コンテナデータ型 — Python 3.7.3 ドキュメント