Day 076-Q1 | Centroid Decomposition + 区間加算クエリ

問題

$N$ 頂点の辺重み付き木が与えられる。各頂点 $v$ に値 $a_v$ が設定されており、以下の $Q$ 個のクエリを処理せよ。

クエリ 1: 1 v x — 頂点 $v$ の値に $x$ を加算する
クエリ 2: 2 v — 頂点 $v$ から距離が $d$ 以下のすべての頂点の値の合計を出力する

距離は辺重みの和として定義する。$d$ は入力で事前に1つ与えられる。

制約

パラメータ	範囲	備考
$N$	$2 \le N \le 10^5$	頂点数
$Q$	$1 \le Q \le 10^5$	クエリ数
$d$	$0 \le d \le 10^{15}$	距離閾値
$w_i$	$1 \le w_i \le 10^9$	辺重み
$a_v, x$	$-10^9 \le a_v, x \le 10^9$	頂点値・加算値

入出力例

入力例1

出力例1

15
25

頂点3から距離3以内: 全頂点(距離0〜2)。初期値和=15。頂点1に+10後=25。

概念図: 重心分解の階層構造

重心分解では各頂点 $v$ から祖先重心 $c$ を $O(\log N)$ 個辿る。各重心 $c$ に距離ソート済み BIT を持たせ、更新・クエリとも $O(\log^2 N)$。

ヒント

ヒント1（方向性）

重心分解（Centroid Decomposition）で木を $O(\log N)$ 深さの階層に分解する。各頂点 $v$ とその祖先重心 $c$ の組に対して距離を事前計算し、BIT を使って「距離 $\le d-\text{dist}(v,c)$ の頂点の値の合計」を効率よく取得する。

ヒント2（アプローチ）

重心分解を行い、各頂点の祖先重心リスト（重心, 距離）を記録する
各重心 $c$ に対して「$c$ からの距離のソート済みリスト + BIT」を用意し初期値を登録
更新: 各祖先重心 $c$ の BIT を dist(v, c) の位置で更新
クエリ: 各祖先重心 $c$ で dist ≤ D - dist(v,c) の範囲の和を取り、包除で重複を除去

ヒント3（ほぼ答え）

# 各祖先重心を辿ってクエリ
def query_correct(v):
    ans = 0
    prev_c = -1
    for c, d in ancestors[v]:
        limit = D - d
        r = bisect_right(cent_sorted_dists[c], max(0, limit)) - 1
        if limit >= 0 and r >= 0:
            ans += cent_bits[c].range_sum(0, r)
        if prev_c != -1:
            for cc, dd in ancestors[v]:
                if cc == prev_c:
                    lim2 = D - dd
                    r2 = bisect_right(cent_sorted_dists[prev_c], max(0, lim2)) - 1
                    if lim2 >= 0 and r2 >= 0:
                        ans -= cent_bits[prev_c].range_sum(0, r2)
                    break
        prev_c = c
    return ans

模範解答

import sys
from collections import defaultdict
from bisect import bisect_right
input = sys.stdin.readline

def main():
    sys.setrecursionlimit(300000)
    N, Q, D = map(int, input().split())
    A = list(map(int, input().split()))
    graph = defaultdict(list)
    for _ in range(N - 1):
        u, v, w = map(int, input().split())
        u -= 1; v -= 1
        graph[u].append((v, w))
        graph[v].append((u, w))

    subtree_sz = [0] * N
    removed = [False] * N
    cent_par = [-1] * N
    cent_dist = [[] for _ in range(N)]

    def calc_size(v, p):
        subtree_sz[v] = 1
        for u, w in graph[v]:
            if u != p and not removed[u]:
                calc_size(u, v)
                subtree_sz[v] += subtree_sz[u]

    def find_centroid(v, p, tree_sz):
        for u, w in graph[v]:
            if u != p and not removed[u]:
                if subtree_sz[u] > tree_sz // 2:
                    return find_centroid(u, v, tree_sz)
        return v

    def collect_dist(v, p, c, d):
        cent_dist[c].append((v, d))
        for u, w in graph[v]:
            if u != p and not removed[u]:
                collect_dist(u, v, c, d + w)

    def decompose(v, par):
        calc_size(v, -1)
        c = find_centroid(v, -1, subtree_sz[v])
        cent_par[c] = par
        collect_dist(c, -1, c, 0)
        removed[c] = True
        for u, w in graph[c]:
            if not removed[u]:
                decompose(u, c)

    decompose(0, -1)

    ancestors = [[] for _ in range(N)]
    for c in range(N):
        for v, d in cent_dist[c]:
            ancestors[v].append((c, d))

    cent_sorted_dists = []
    for c in range(N):
        dists = sorted(d for v, d in cent_dist[c])
        cent_sorted_dists.append(dists)

    class BIT:
        def __init__(self, n):
            self.n = n
            self.bit = [0] * (n + 1)
        def add(self, i, x):
            i += 1
            while i <= self.n:
                self.bit[i] += x
                i += i & (-i)
        def sum(self, i):
            i += 1; s = 0
            while i > 0:
                s += self.bit[i]; i -= i & (-i)
            return s
        def range_sum(self, l, r):
            if l > r: return 0
            return self.sum(r) - (self.sum(l - 1) if l > 0 else 0)

    cent_bits = [BIT(len(cent_sorted_dists[c])) for c in range(N)]
    for v in range(N):
        for c, d in ancestors[v]:
            idx = bisect_right(cent_sorted_dists[c], d) - 1
            cent_bits[c].add(idx, A[v])

    def update(v, x):
        for c, d in ancestors[v]:
            idx = bisect_right(cent_sorted_dists[c], d) - 1
            cent_bits[c].add(idx, x)

    def query_correct(v):
        ans = 0
        prev_c = -1
        for c, d in ancestors[v]:
            limit = D - d
            r = bisect_right(cent_sorted_dists[c], max(0, limit)) - 1
            if limit >= 0 and r >= 0:
                ans += cent_bits[c].range_sum(0, r)
            if prev_c != -1:
                for cc, dd in ancestors[v]:
                    if cc == prev_c:
                        lim2 = D - dd
                        r2 = bisect_right(cent_sorted_dists[prev_c], max(0, lim2)) - 1
                        if lim2 >= 0 and r2 >= 0:
                            ans -= cent_bits[prev_c].range_sum(0, r2)
                        break
            prev_c = c
        return ans

    out = []
    for _ in range(Q):
        line = input().split()
        if line[0] == '1':
            v, x = int(line[1]) - 1, int(line[2])
            update(v, x)
        else:
            v = int(line[1]) - 1
            out.append(query_correct(v))
    print('\n'.join(map(str, out)))

main()

Step-by-Step 解説

Step 1: 重心とは何か

木の重心は「取り除くと各連結成分のサイズが元の木の半分以下になる頂点」。常に存在し、$O(N)$ で見つけられる。

Step 2: 重心分解の再帰構造

重心を取り除き残りの各連結成分でも再帰的に重心を求める。深さは $O(\log N)$（各ステップで問題サイズが半分以下）。各頂点は階層上 $O(\log N)$ 個の重心の「管轄」にある。

Step 3: 距離 BIT の構築

各重心 $c$ について、管轄する全頂点の距離をソートしてリストを作り、BIT を乗せる。「距離 $\le X$ の頂点の値の合計」は bisect + BIT.range_sum で $O(\log N)$。

Step 4: 更新と重複除去

頂点 $v$ を更新する時、各祖先重心 $c$ の BIT を更新。クエリ時は各重心でカウントするが同一頂点が複数回計上されないよう、親重心分を包除で引く。

Step 5: 計算量

各クエリ: $O(\log N \cdot \log N) = O(\log^2 N)$。全体: $O((N + Q) \log^2 N)$。

よくあるミス

ミス	原因	正しい書き方
重複カウント	包除が不完全	各レベルで親重心分を必ず引く
距離計算の誤り	重心からの距離を正しく記録していない	`collect_dist` 時に累積距離を渡す
removed 配列の忘れ	分解後の頂点を除外しない	`removed[c] = True` を忘れずに

次のステップ

発展問題: 「木上の頂点のうち距離 $d$ 以内にある頂点数が最も多い頂点を求めよ」→ 同様の重心分解 + ソートで解ける

自己評価

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

自分の回答:

気づき・メモ: