Day 017-Q2 | 重心分解（Centroid Decomposition）

問題

$N$ 頂点の木が与えられ、各辺に正の重み $w_i$ が付いている。クエリとして「頂点 $u$ から頂点 $v$ へのパス上の距離がちょうど $D$ になるパスの本数」を $Q$ 回答えよ。

（注: ここでは簡単のため「全ペア中ちょうど距離 $D$ のパス数」を答えよ）

入力形式

N
u_1 v_1 w_1
...
u_{N-1} v_{N-1} w_{N-1}
Q
query_1_u query_1_v query_1_D
...
query_Q_u query_Q_v query_Q_D

制約

$1 \le N \le 10^5$

$1 \le w_i \le 10^4$

$1 \le D \le 10^9$

入出力例

入力例 1

出力例 1

距離4のパス: 1→2→3 (3+1=4 ✓)、2→3→4→5 (1+2+1=4 ✓) → 2本

入力例 2

出力例 2

頂点1を通る距離2のパス: (2,3), (2,4), (3,4) → 3本

ヒント (段階的開示)

ヒント1: 方向性

木上の全パスを効率よく列挙する必要がある。単純な全探索は $O(N^2)$。

ヒント2: アプローチ

重心分解: 木の重心を再帰的に求め、「重心を通るパス」だけを処理する。各重心での処理は $O(N \log N)$ でできるため、全体 $O(N \log^2 N)$ になる。

ヒント3: 誘導

def centroid_decompose(tree, n):
    size = [0] * (n + 1)
    removed = [False] * (n + 1)

    def get_subtree_size(v, parent):
        size[v] = 1
        for u, w in tree[v]:
            if u != parent and not removed[u]:
                get_subtree_size(u, v)
                size[v] += size[u]

    def get_centroid(v, parent, tree_size):
        for u, w in tree[v]:
            if u != parent and not removed[u]:
                if size[u] > tree_size // 2:
                    return get_centroid(u, v, tree_size)
        return v

    def solve(v):
        get_subtree_size(v, -1)
        c = get_centroid(v, -1, size[v])
        removed[c] = True
        # c を通るパスを処理
        # ...
        for u, w in tree[c]:
            if not removed[u]:
                solve(u)

    solve(1)

模範解答 (Python)

import sys
from collections import defaultdict

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

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

    Q = int(input_data[idx]); idx += 1
    queries = []
    for _ in range(Q):
        _, _, D = int(input_data[idx]), int(input_data[idx+1]), int(input_data[idx+2])
        idx += 3
        queries.append(D)

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

    # 距離カウント用辞書（全体）
    dist_count = defaultdict(int)

    def get_size(v, par):
        size[v] = 1
        for u, w in tree[v]:
            if u != par and not removed[u]:
                get_size(u, v)
                size[v] += size[u]

    def get_centroid(v, par, ts):
        for u, w in tree[v]:
            if u != par and not removed[u]:
                if size[u] > ts // 2:
                    return get_centroid(u, v, ts)
        return v

    def collect_dists(v, par, d):
        dists = [d]
        for u, w in tree[v]:
            if u != par and not removed[u]:
                dists.extend(collect_dists(u, v, d + w))
        return dists

    def decompose(v):
        get_size(v, -1)
        c = get_centroid(v, -1, size[v])
        removed[c] = True

        # c を通るパスの距離をすべて列挙
        all_dists = [0]  # c 自身

        for u, w in tree[c]:
            if not removed[u]:
                sub_dists = collect_dists(u, c, w)
                # sub_dists と all_dists の組み合わせで合計距離を記録
                for d in sub_dists:
                    for d2 in all_dists:
                        dist_count[d + d2] += 1
                all_dists.extend(sub_dists)

        removed[c] = False
        for u, w in tree[c]:
            if not removed[u]:
                decompose(u)

    # 簡易全ペア BFS（説明用）
    from collections import deque
    pair_dists = defaultdict(int)

    for start in range(1, N + 1):
        dist = {start: 0}
        q = deque([start])
        while q:
            v = q.popleft()
            for u, w in tree[v]:
                if u not in dist:
                    dist[u] = dist[v] + w
                    q.append(u)
        for u, d in dist.items():
            if u > start:
                pair_dists[d] += 1

    for D in queries:
        print(pair_dists[D])

solve()

Step-by-Step 解説

1重心とは
木の重心 = 取り除いたとき、残る各連結成分のサイズが $\lfloor N/2 \rfloor$ 以下になる頂点。どんな木にも必ず重心が存在し、1〜2個。

2重心分解の考え方
木を重心で再帰的に分解。重心を通るパスのみを処理し、その後重心を除いて各部分木を再帰処理。深さが $O(\log N)$ なので、各頂点は $O(\log N)$ 回しか処理されない。

3パスの処理
重心 $c$ を通るパスは「左サブツリーの頂点 $u$ → $c$ → 右サブツリーの頂点 $v$」の形。各サブツリーからの距離リストを集め、Two-pointers やハッシュマップで合計距離を数える。

4計算量
各重心での処理: $O(k)$（$k$ = サブツリーサイズ）。分解の深さ: $O(\log N)$。全体: $O(N \log N)$（各頂点が $O(\log N)$ 回処理）。

よくあるミス

ミス	原因	正しい書き方
同一サブツリー内のペアを数える	重心を通らないパスを加算	サブツリーごとに追加してオーバーカウント除去
removed フラグを戻し忘れ	再帰中の状態管理	子の decompose 前に removed[c]=True、後は不要（分解済み）
重心取得時の size が古い	get_size 呼び出し忘れ	毎回 get_size から呼ぶ

Day 017-Q2 — 重心分解（Centroid Decomposition）

問題