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

問題

$N$頂点の木が与えられる。辺$i$は頂点$u_i,v_i$を結び重み$w_i$を持つ。2頂点$u,v$の距離$\mathrm{dist}(u,v)$をパス上の辺重みの総和とする。整数$K$が与えられるので、$\mathrm{dist}(u,v)\le K$を満たす頂点対$(u,v)$（$u

重心分解は、木を「取り除くと残りの各連結成分のサイズが元の半分以下になる頂点＝重心」で分割し、「重心を経由するパス」をまとめて処理してから残りの部分木に再帰する分割統治法。深さは$O(\log N)$で、全体$O(N\log^2 N)$で全パスの統計を計算できる。

入力形式

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

制約

$2 \le N \le 2\times10^5$

$1 \le K \le 10^{14}$

$1 \le w_i \le 10^9$

与えられるグラフは木

入出力例

入力例1

出力例1

距離は $d(1,2)=1,d(1,3)=3,d(1,4)=6,d(2,3)=2,d(2,4)=5,d(3,4)=3$。3以下は$(1,2),(1,3),(2,3),(3,4)$の4組。

概念図: 重心を通るパスの数え上げ

ヒント（段階的開示）

ヒント1: 方向性

木上のすべてのパスは、適切な1頂点で分割すれば「その頂点を通るパス」と「それを含まないより小さい部分木内で完結するパス」に分けられる。この頂点として重心を選ぶと、部分木サイズが毎回半分以下になり再帰の深さが$O(\log N)$に収まる。

ヒント2: アプローチ

重心$c$を削除し、$c$から各頂点への距離をDFSで求める。ソートしてtwo-pointerで「距離の和が$K$以下」のペア数を数える。ただしこれは同じ子部分木内のペアまで含んでしまうので、各子部分木ごとに同じ計算をして引き算する。その後$c$を除いた各部分木に再帰する。

ヒント3: 誘導（コード骨格）

def count_pairs_le_k(dists, K):
    dists.sort()
    i, j = 0, len(dists) - 1
    cnt = 0
    while i < j:
        if dists[i] + dists[j] <= K:
            cnt += j - i
            i += 1
        else:
            j -= 1
    return cnt

def decompose(v):
    c = 重心(v)
    削除(c)
    all_dists = [0] + (cから各部分木への距離すべて)
    ans += count_pairs_le_k(all_dists, K)
    for 各子部分木 sub:
        ans -= count_pairs_le_k(sub の距離リスト, K)
    for 各子部分木 sub:
        decompose(sub)

模範解答 (Python)

import sys


def solve():
    data = sys.stdin.buffer.read().split()
    idx = 0
    n = int(data[idx]); idx += 1
    k = int(data[idx]); idx += 1
    graph = [[] for _ in range(n + 1)]
    for _ in range(n - 1):
        u = int(data[idx]); v = int(data[idx + 1]); w = int(data[idx + 2])
        idx += 3
        graph[u].append((v, w))
        graph[v].append((u, w))

    sys.setrecursionlimit(500000)

    removed = [False] * (n + 1)
    subtree_size = [0] * (n + 1)

    def calc_size(v, parent):
        subtree_size[v] = 1
        for to, w in graph[v]:
            if to != parent and not removed[to]:
                calc_size(to, v)
                subtree_size[v] += subtree_size[to]
        return subtree_size[v]

    def find_centroid(v, parent, tree_size):
        for to, w in graph[v]:
            if to != parent and not removed[to] and subtree_size[to] > tree_size // 2:
                return find_centroid(to, v, tree_size)
        return v

    def collect_dists(v, parent, dist, out):
        out.append(dist)
        for to, w in graph[v]:
            if to != parent and not removed[to]:
                collect_dists(to, v, dist + w, out)

    def count_pairs_le_k(dists):
        dists.sort()
        i, j = 0, len(dists) - 1
        cnt = 0
        while i < j:
            if dists[i] + dists[j] <= k:
                cnt += j - i
                i += 1
            else:
                j -= 1
        return cnt

    ans = 0

    def decompose(v):
        nonlocal ans
        tree_size = calc_size(v, -1)
        centroid = find_centroid(v, -1, tree_size)
        removed[centroid] = True

        neighbors = [(to, w) for to, w in graph[centroid] if not removed[to]]

        all_dists = [0]
        child_dist_lists = []
        for to, w in neighbors:
            sub = []
            collect_dists(to, centroid, w, sub)
            child_dist_lists.append(sub)
            all_dists.extend(sub)

        ans += count_pairs_le_k(all_dists)
        for sub in child_dist_lists:
            ans -= count_pairs_le_k(sub)

        for to, w in neighbors:
            decompose(to)

    decompose(1)
    print(ans)


solve()

計算量: $O(N\log^2 N)$（重心分解 $O(\log N)$ 層 × 各層 $O(N\log N)$ のソート）。ランダム木300ケースで愚直全点対計算と一致することを確認済み。

Step-by-Step 解説

1重心を求める
DFSで部分木サイズを求め（`calc_size`）、「取り除くとどの方向の連結成分もサイズが半分以下」になる頂点を`find_centroid`で見つける。

2重心を通るパスをまとめて数える
重心$c$から到達可能な全頂点への距離を集め（$c$自身は距離0）、ソートしてtwo-pointerで距離の和が$K$以下のペア数を数える。

3同一部分木内の過大カウントを補正
$c$の各子部分木ごとに同じ計算をして、$c$を経由しない同一部分木内ペアの分を引く。

4再帰
$c$を削除済みにし、各子部分木に同じ手順を再帰適用。重心分解の性質上、深さは$O(\log N)$。

よくあるミス

ミス	原因	正しい書き方
重心を求める際に古いサイズを使い続ける	再帰のたびにその連結成分でサイズ再計算が必要	`decompose(v)`の最初に必ず`calc_size`をやり直す
同じ部分木内のペアを引き算し忘れる	重心を通らないペアまで過大カウントする	子部分木ごとに`count_pairs_le_k`を計算し全体から引く
`removed`フラグの更新を忘れる	削除済み重心を再訪問し二重カウント・無限ループの原因	重心を選んだら即座に`removed[centroid]=True`
two-pointerで`j-i`個まとめて数えるのを忘れる	1組ずつ+1しようとして計算量が悪化・実装も煩雑	`dists[i]+dists[j]<=K`成立時は`j-i`個すべて条件を満たすことを利用

次のステップ

発展: 「距離がちょうど$K$」のパスを数える古典問題（Codeforces 161D相当）との関係を確認する
発展: 頂点重みを付け「距離が$K$以下のペアの重みの積の総和」など集計内容を一般化する
次回予告: 線形漸化式の高速計算（行列累乗・Companion Matrix）— $N\le10^{18}$項目を高速に求める

自己評価

完全理解おおむね理解要復習

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

問題