木の直径

適当な頂点から最も遠い頂点 $u$ を求め、さらに $u$ から最も遠い頂点 $v$ を求めると、 $u$ と $v$ を結ぶパスが直径になります。BFS で書いているので再帰は使いません。

• tree_diameter(g): 重みなしの木 g[v] = [to, ...] について、(直径（辺の本数）, 直径をなすパスの頂点列) を返す。 $O(N)$
• tree_diameter_weighted(g): 重み付きの木 g[v] = [(to, w), ...]（ $w \ge 0$）について、(直径, パスの頂点列) を返す。 $O(N)$

コード

def tree_diameter_weighted(g):
    # g[v] = [(to, w), ...]、w は非負
    n = len(g)

    def bfs(s):
        dist = [-1] * n
        par = [-1] * n
        dist[s] = 0
        order = [s]
        for v in order:
            for to, w in g[v]:
                if dist[to] == -1:
                    dist[to] = dist[v] + w
                    par[to] = v
                    order.append(to)
        return dist, par

    dist, _ = bfs(0)
    u = dist.index(max(dist))
    dist, par = bfs(u)
    v = dist.index(max(dist))
    path = [v]
    while path[-1] != u:
        path.append(par[path[-1]])
    return dist[v], path


def tree_diameter(g):
    # g[v] = [to, ...]（重みなし）
    return tree_diameter_weighted([[(to, 1) for to in gv] for gv in g])