Day 075-Q3 — Rerooting DP（全方位木DP・全頂点深さ総和 O(N)）

問題

$N$ 頂点の木が与えられる。各頂点 $v$ に対して、「$v$ を根としたとき、各頂点の深さの総和」$S_v = \sum_{u=0}^{N-1} \text{depth}_v(u)$ を全頂点について求めよ。

入出力例

5
0 1
0 2
1 3
1 4

6 5 10 8 8

$S_0 = 0+1+1+2+2 = 6$。$S_1 = 1+0+2+1+1 = 5$。$S_2 = 1+2+0+3+3 = 9$（木の構造による）。

概念図: Rerooting の遷移

根を $p$ から $v$（部分木サイズ $\text{sz}[v]$）に移すと、$v$ の部分木の $\text{sz}[v]$ 頂点は深さが 1 減り、それ以外の $N - \text{sz}[v]$ 頂点は深さが 1 増える。

ヒント

# 上向き伝播
for v in order:  # BFS 順（pre-order）
    for u in adj[v]:
        if u != parent[v]:
            # u を根にすると:
            # - u の部分木 sz[u] 頂点: 深さ -1 (v→u の辺が根方向になる)
            # - 残り N-sz[u] 頂点: 深さ +1
            S[u] = S[v] + N - 2 * sz[u]

模範解答

import sys
input = sys.stdin.readline

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

    root = 0
    sz = [1] * N; dp = [0] * N
    parent = [-1] * N; order = []

    # BFS で pre-order 列を作成（DFS も可）
    from collections import deque
    q = deque([root]); parent[root] = -1
    while q:
        v = q.popleft(); order.append(v)
        for u in adj[v]:
            if u != parent[v]:
                parent[u] = v; q.append(u)

    # post-order で下向き DP
    for v in reversed(order):
        p = parent[v]
        if p != -1:
            dp[p] += dp[v] + sz[v]
            sz[p] += sz[v]

    # 上向き伝播
    S = [0] * N; S[root] = dp[root]
    for v in order:
        for u in adj[v]:
            if u != parent[v]:
                S[u] = S[v] + N - 2 * sz[u]

    print(*S)

solve()

Step-by-Step 解説

Step 1: 下向き DP の直観

頂点 0 を根として、$\text{dp}[v]$ = 「頂点 $v$ の部分木内の頂点について、$v$ からの相対深さの総和」を post-order で計算する。子 $c$ の部分木の全頂点は $v$ から見ると深さが 1 増えるので:

$$\text{dp}[v] \mathrel{+}= \text{dp}[c] + \text{sz}[c]$$

Step 2: 根の移し替えの遷移式

根を $p$ から子 $v$ に移すと（$\text{sz}[v]$ は旧根基準の $v$ の部分木サイズ）:

• $v$ の部分木の $\text{sz}[v]$ 頂点: 全員の深さが 1 減る
• それ以外の $N - \text{sz}[v]$ 頂点: 全員の深さが 1 増える

$$S[v] = S[p] + (N - \text{sz}[v]) - \text{sz}[v] = S[p] + N - 2 \cdot \text{sz}[v]$$

Step 3: BFS 順（pre-order）で伝播

BFS で得た order はそのまま pre-order。親 $p$ の $S[p]$ が確定してから子 $v$ の $S[v]$ を計算するため、order の順序のまま処理すればよい。

Step 4: 計算量

BFS $O(N)$ + post-order 集約 $O(N)$ + 伝播 $O(N)$ で全体 $O(N)$。

よくあるミス

次のステップ

• 発展問題: 辺重みが存在する場合の全頂点距離和、またはパス上の辺重み最大値の全頂点版（部分木マージで拡張）

自己評価