Day 023-Q3 | Rerooting (全方位木DP)

問題

$N$ 頂点の木が与えられる。各頂点 $v$ について「$v$ を根としたとき、全頂点間の距離の総和 $\sum_{u} \text{dist}(v, u)$」を求めよ。

入力形式

N
u_1 v_1
...
u_{N-1} v_{N-1}

制約

$2 \le N \le 2 \times 10^5$

辺は重みなし

木は連結

入出力例

入力例 1

出力例 1

10 7 6 7 10

ヒント (段階的開示)

ヒント1: 方向性

全頂点を根にした答えを素朴に求めると $O(N^2)$。Rerooting で $O(N)$ に。

ヒント2: アプローチ

まず根 0 で部分木サイズ sub_size と「部分木内の距離和」down を計算（下り）。次に親→子へ伝播（上り）して全頂点の答えを得る。

ヒント3: 誘導

# 子 c に根を移すと:
# c 側の sub_size[c] 頂点は 1 近くなり
# 残り N - sub_size[c] 頂点は 1 遠くなる
ans[c] = ans[v] - sub_size[c] + (N - sub_size[c])

模範解答 (Python)

import sys
from collections import defaultdict, deque
input = sys.stdin.readline

def solve():
    N = int(input())
    graph = defaultdict(list)
    for _ in range(N - 1):
        u, v = map(int, input().split())
        graph[u].append(v)
        graph[v].append(u)

    sub_size = [1] * N
    down = [0] * N
    order = []
    parent = [-1] * N

    visited = [False] * N
    q = deque([0])
    visited[0] = True
    while q:
        v = q.popleft()
        order.append(v)
        for u in graph[v]:
            if not visited[u]:
                visited[u] = True
                parent[u] = v
                q.append(u)

    for v in reversed(order):
        for u in graph[v]:
            if u == parent[v]:
                continue
            sub_size[v] += sub_size[u]
            down[v] += down[u] + sub_size[u]

    ans = [0] * N
    ans[0] = down[0]
    for v in order:
        for u in graph[v]:
            if u == parent[v]:
                continue
            ans[u] = ans[v] - sub_size[u] + (N - sub_size[u])

    print(*ans)

solve()

Step-by-Step 解説

1問題の変換
根 0 → 子 c に変えると、c 部分木は 1 減、それ以外は 1 増。ans[c] = ans[v] - sub_size[c] + (N - sub_size[c])

2下りDP
葉から根への BFS 逆順。down[v] += down[u] + sub_size[u]。

3上りDP
根から葉への BFS 正順。各子について Rerooting 式を適用。$O(N)$。

4再帰を避ける
$N = 2 \times 10^5$ の竹で再帰深度が $N$ に達するため、BFS でトポロジカル順を求めて反復処理。

よくあるミス

ミス	原因	正しい書き方
再帰でスタックオーバーフロー	深い木で再帰	BFS + イテレーティブ DP
親を子として処理	親判定漏れ	`if u == parent[v]: continue`
ans[c] の式ミス	`sub_size[c]` の役割誤解	`-sub_size[c] + (N - sub_size[c])`

Day 023-Q3 — Rerooting (全方位木DP)

問題