GRL_5_B - Height of a Tree / PythonでAizu Online Judgeを解く

Posted on 2019/06/30

木の高さ

非負の重みを持つ無向辺からなる木 \(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

制約

  • \(s_i\), \(t_i\) は \(i\) 番目の辺の端点
  • \(w_i\) は \(i\) 番目の辺の重み
  • 木の頂点には、それぞれ \(0\) から \(n-1\) の番号が付けられているものとします
  • \(1 \leq n \leq 10,000\)
  • \(0 \leq w_i \leq 1,000\)

サンプル入力

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

考え方

とりあえず全探索から考える 🤔

ある頂点を根として一番遠い頂点を求めるだけであればその頂点を起点に深さ優先探索などをすれば良くて高々 \(\mathcal{O(n) }\) に収まりますが、全ての頂点についてそれをやろうと \(\mathcal{O(n^2)}\) になるのでどうしたものかという感じです。

唐突に「木の直径」を思い出す

木の直径とは「最も離れている頂点間の距離を指すもの」でした。木の直径の端点ではないような頂点までの距離が木の高さになるのであればその頂点を使ってより大きな木の直径が求まるはずで定義と矛盾します。つまりある頂点の高さはその木の直径となるような2頂点いずれかとの距離になるはずです。

実装

実装については木の直径(GRL_5_A)で使ったコードをほぼ流用することができます。直径の端点となる2点についてそれぞれを起点に各頂点までの距離を求め、各頂点について2点との距離の最大値を求めてやると高さが求まります。

%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, best
    
    def run(root):
        heappush(queue, (0, root))
        return find(root)
    
    ret0, best0 = run(0)
    ret1, best1 = run(ret0[1])
    ret2, best2 = run(ret1[1])
    
    for i in range(n):
        yield max(best1[i], best2[i])

サンプル出力

%python
with open('/tmp/input1.txt') as f:
    for d in solve(*input(f)):
        print(d)
5
3
4
5

証明は背理法でいいのだろうか?🤔