Day 086-Q3 — Heavy-Light Decomposition（パス区間和・点更新・$O(\log^2 N)$）

問題

$N$ 頂点の木（頂点 $0$ を根）があり、頂点 $i$ は値 $a_i$ を持つ。$Q$ 個のクエリを処理せよ。

• 1 u x: 頂点 $u$ の値に $x$ を加算する。
• 2 u v: $u$–$v$ パス上（両端含む）の頂点値の総和を答えよ。

入出力例

3
1 2 3
0 1
1 2
3
2 0 2
1 1 5
2 0 2

6
11

初期パス $0\text{-}1\text{-}2 = 1+2+3 = 6$。頂点1に5加算後 $1+7+3 = 11$。

概念図: heavy 辺によるチェーン分解

ヒント

while head[u] != head[v]:
    if depth[head[u]] < depth[head[v]]: u, v = v, u
    res += range_sum(pos[head[u]], pos[u])
    u = parent[head[u]]
if depth[u] > depth[v]: u, v = v, u
res += range_sum(pos[u], pos[v])

模範解答

import sys

def solve():
    data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(data[idx]); idx += 1
    val = [int(data[idx + i]) for i in range(N)]; idx += N
    g = [[] for _ in range(N)]
    for _ in range(N - 1):
        u = int(data[idx]); v = int(data[idx + 1]); idx += 2
        g[u].append(v); g[v].append(u)

    parent = [-1] * N; depth = [0] * N; size = [1] * N; order = []
    visited = [False] * N; visited[0] = True; st = [0]
    while st:
        u = st.pop(); order.append(u)
        for w in g[u]:
            if not visited[w]:
                visited[w] = True; parent[w] = u; depth[w] = depth[u] + 1; st.append(w)
    for u in reversed(order):
        if parent[u] != -1:
            size[parent[u]] += size[u]
    heavy = [-1] * N
    for u in order:
        mx = 0
        for w in g[u]:
            if w != parent[u] and size[w] > mx:
                mx = size[w]; heavy[u] = w

    head = [0] * N; pos = [0] * N; cur = 0
    for u in order:
        if parent[u] == -1 or heavy[parent[u]] != u:
            v = u
            while v != -1:
                head[v] = u; pos[v] = cur; cur += 1
                v = heavy[v]

    bit = [0] * (N + 1)
    def upd(i, x):
        i += 1
        while i <= N:
            bit[i] += x; i += i & -i
    def pref(i):  # sum of [0, i)
        s = 0
        while i > 0:
            s += bit[i]; i -= i & -i
        return s
    for i in range(N):
        upd(pos[i], val[i])

    Q = int(data[idx]); idx += 1
    out = []
    for _ in range(Q):
        t = int(data[idx]); idx += 1
        if t == 1:
            u = int(data[idx]); x = int(data[idx + 1]); idx += 2
            upd(pos[u], x)
        else:
            u = int(data[idx]); v = int(data[idx + 1]); idx += 2
            res = 0
            while head[u] != head[v]:
                if depth[head[u]] < depth[head[v]]:
                    u, v = v, u
                res += pref(pos[u] + 1) - pref(pos[head[u]])
                u = parent[head[u]]
            if depth[u] > depth[v]:
                u, v = v, u
            res += pref(pos[v] + 1) - pref(pos[u])
            out.append(res)
    print('\n'.join(map(str, out)))

solve()

計算量: 前処理 $O(N)$、各クエリ $O(\log^2 N)$（チェーン $O(\log N)$ × BIT $O(\log N)$）。

Step-by-Step 解説

Step 1: DFS でサイズ・親・深さ

反復 DFS で order を得て逆順で部分木サイズを集計（再帰上限回避）。

Step 2: heavy child の決定

各頂点で部分木サイズ最大の子を heavy とする。heavy 辺を辿ると 1 本のチェーン。

Step 3: チェーン分解と pos 割り当て

チェーン先頭から heavy を辿り連続番号を振る。同一チェーンは pos が連続。

Step 4: パスクエリ

よくあるミス

次のステップ

• 発展問題: 辺重みパス最大値（辺を子頂点に載せる + セグメント木 max）
• 発展問題: パス区間加算 + パス区間和（HLD + 遅延伝播セグメント木）

自己評価