Day 033-Q5 | 木の重心分解 + 動的計画法

問題

$N$ 頂点の重み付き木で、距離がちょうど $K$ である頂点対 $(u, v)$（$u < v$）の数を求めよ。距離は辺の重みの和。

入力形式

N K
u_1 v_1 w_1
...
u_{N-1} v_{N-1} w_{N-1}

制約

$2 \le N \le 10^5$

$1 \le K \le 10^9$

$1 \le w_i \le 10^4$

入出力例

入力例 1

出力例 1

ヒント (段階的開示)

ヒント1: 方向性

全頂点対の距離 = $K$ を数える。重心分解（Centroid Decomposition）で $O(N \log^2 N)$。

ヒント2: アプローチ

重心 $c$ を求める → 重心経由パスをカウント → 重心削除 → 再帰。各頂点から $c$ への距離 $d_v$ で $d_u + d_v = K$ のペアを集計。

ヒント3: 誘導

同一部分木内の二重カウント回避は「先にマッチング → 後に追加」の順序で。

模範解答 (Python)

import sys
from collections import defaultdict

def solve():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N, K = int(input_data[idx]), int(input_data[idx+1]); idx += 2

    adj = [[] for _ in range(N + 1)]
    for _ in range(N - 1):
        u, v, w = int(input_data[idx]), int(input_data[idx+1]), int(input_data[idx+2]); idx += 3
        adj[u].append((v, w))
        adj[v].append((u, w))

    removed = [False] * (N + 1)
    sub_size = [0] * (N + 1)
    answer = 0

    sys.setrecursionlimit(5 * 10**5)

    def calc_size(v, p):
        sub_size[v] = 1
        for u, _ in adj[v]:
            if u != p and not removed[u]:
                calc_size(u, v)
                sub_size[v] += sub_size[u]

    def find_centroid(v, p, tree_size):
        for u, _ in adj[v]:
            if u != p and not removed[u]:
                if sub_size[u] > tree_size // 2:
                    return find_centroid(u, v, tree_size)
        return v

    def get_dists(v, p, d, result):
        result.append(d)
        for u, w in adj[v]:
            if u != p and not removed[u]:
                get_dists(u, v, d + w, result)

    def process(v):
        nonlocal answer
        calc_size(v, -1)
        c = find_centroid(v, -1, sub_size[v])

        dist_count = defaultdict(int)
        dist_count[0] = 1

        for u, w in adj[c]:
            if removed[u]:
                continue
            dists = []
            get_dists(u, c, w, dists)

            for d in dists:
                need = K - d
                if need >= 0:
                    answer += dist_count[need]

            for d in dists:
                dist_count[d] += 1

        removed[c] = True
        for u, _ in adj[c]:
            if not removed[u]:
                process(u)

    process(1)
    print(answer)

solve()

Step-by-Step 解説

1重心分解の原理
重心 $c$: 削除時に生じる連結成分のサイズが $\le N/2$。任意の木に必ず存在し $\le 2$ 個。分割統治の深さ $O(\log N)$。

2重心通過パスをカウント
各頂点 $v$ から重心 $c$ への距離 $d_v$ を列挙。$d_u + d_v = K$ を defaultdict でカウント。

3二重カウント回避
各部分木について「先にマッチング → 後に dist_count に追加」の順序で、同一部分木内パスを除外。

4再帰
重心削除後、各連結成分に対して再帰的に process。

計算量

重心分解の深さ: $O(\log N)$
各レベルの処理: $O(N)$
全体: $O(N \log N)$（hash 衝突を含めて実質 $O(N \log N)$）

よくあるミス

ミス	原因	正しい書き方
同一部分木内パスを二重カウント	追加後にマッチング	先にマッチング、後に追加
`removed` 管理漏れ	削除重心を再訪	すべての関数で `removed` 確認
再帰スタック超過	$N = 10^5$ で線形深さ	`sys.setrecursionlimit(5*10^5)`
$K - d < 0$ を未考慮	距離は非負	`if need >= 0` でチェック

Day 033-Q5 — 木の重心分解 + 動的計画法（Centroid Decomp + Path DP）

問題