GRL_5_A - Diameter of a Tree / PythonでAizu Online Judgeを解く

Posted on 2019/06/23

木の直径

非負の重みを持つ無向辺からなる木 \(T\) の直径を求めてください。ここで木の直径とは最も離れている頂点間の距離を指すものとします。

入力

%python
def input(f):
    n, = map(int, f.readline().split())
    s = [None for _ in range(n-1)]
    t = [None for _ in range(n-1)]
    w = [None for _ in range(n-1)]
    for i in range(n-1):
        s[i], t[i], w[i] = map(int, f.readline().split())
    return n, s, t, w

制約

  • \(1 \leq n \leq 100,000\)
  • \(0 \leq w \leq 1,000\)

サンプル入力

%sh
cat << EOF > /tmp/input1
4
0 1 2
1 2 1
1 3 3
EOF

%python
with open('/tmp/input1') as f:
    print(input(f))
(4, [0, 1, 1], [1, 2, 3], [2, 1, 3])

考え方

木とは

\(N\) 個の頂点からなるグラフ \(G\) が木のとき以下のような性質が成り立ちます:

  • グラフ \(G\) は閉路を持たない
  • グラフ \(G\) は \(N-1\) 本の辺を持つ

これらの性質があるため深さ優先探索などを活用しやすかったりします。

木の直径を測る

木の直径は以下2つのステップで求まります:

  • 適当な頂点を起点に一番遠い頂点 \(A\) を求める
  • 一番遠い頂点 \(A\) を起点にそこから一番遠い頂点 \(B\) を求める

求めた頂点 \(A\) と頂点 \(B\) の距離が木の直径です。

※ 図のグラフは辺の重みを考慮していません

実装

%python
from heapq import heappush, heappop

def solve(n, s, t, w):
    graph = [[] for _ in range(n)]
    def add_edge(src, dst, w):
        graph[src].append((dst, w))
        graph[dst].append((src, w))
    for i in range(n-1):
        add_edge(s[i], t[i], w[i])
    
    queue = []
    
    def find(root):
        INF = 100000 * 1000 + 1;
        best = [INF for _ in range(n)]
        best[root] = 0
        
        def find_next(cur):
            for nxt_v, nxt_w in graph[cur]:
                dist = best[cur] + nxt_w
                if dist < best[nxt_v]:
                    best[nxt_v] = dist
                    yield dist, nxt_v
        
        def dequeue():
            _, cur_v = heappop(queue)
            for nxt in find_next(cur_v):
                heappush(queue, nxt)
        
        while queue:
            dequeue()

        res = (0, root)
        for i in range(n):
            if best[i] > res[0]:
                res = (best[i], i)
        return res
    
    def run(root):
        heappush(queue, (0, root))
        return find(root)
    
    ret = run(0)
    res = run(ret[1])
    
    return res[0]

アニメーション

%sh
curl -s https://judgedat.u-aizu.ac.jp/testcases/GRL_5_A/7/in > /tmp/input7
head /tmp/input7
20
0 1 13
0 2 3
1 3 2
1 4 9
2 5 0
2 6 8
3 7 12
3 8 16
4 9 16

一番遠いところを見るのはすぐ思いつきそうですが、一番遠いところから一番遠いところを考えるに至るのは難易度が高そうな気がします。