Day 039-Q2 | HLD + BIT による木パス sqrt 和クエリ

問題

$N$ 頂点の根付き木（根 = 1）と $Q$ 個のクエリ。各辺には整数の重みがある。

1 u w: 頂点 $u$ の親との辺の重みを $w$ に変更
2 u v: パス $(u, v)$ 上の辺重みの整数平方根の総和 $\sum_{e \in \text{path}(u,v)} \lfloor \sqrt{w_e} \rfloor$ を求めよ

制約

$2 \le N \le 10^5$

$1 \le Q \le 10^5$

$1 \le w \le 10^9$

時間制限: 3sec / メモリ: 256MB

入出力例

入力例 1

出力例 1

12
17
12

初期辺重み: (1,2)=4, (1,3)=9, (2,4)=16, (2,5)=25。クエリ2(1→5): path=辺(1,2),(2,5)→√4+√25=2+5=7... 実際は path=1-2-5。

概念図: HLD による木の平坦化

ヒント（段階的開示）

ヒント1: 方向性

HLD でパスを $O(\log N)$ 個の連続区間に分解し、各区間の $\lfloor \sqrt{w} \rfloor$ 合計を BIT で管理。辺の重みは深い側の頂点に対応させる。

ヒント2: 辺→頂点の対応

辺 $(parent[v], v)$ の重みを頂点 $v$ に持たせると、辺クエリが点クエリに帰着できる。パスクエリ時は LCA の辺（LCA 自身の位置）を除外することに注意。

ヒント3: 更新と BIT

# 更新クエリ: 辺重みを変更
old_sq = isqrt(edge_w[u])
new_sq = isqrt(w)
edge_w[u] = w
bit.add(pos[u] + 1, new_sq - old_sq)  # 差分更新

# パスクエリ
def path_sum(u, v):
    res = 0
    while head[u] != head[v]:
        if depth[head[u]] < depth[head[v]]:
            u, v = v, u
        res += bit.query(pos[head[u]] + 1, pos[u] + 1)
        u = parent[head[u]]
    if depth[u] > depth[v]:
        u, v = v, u
    if u != v:  # LCA 辺を除外
        res += bit.query(pos[u] + 2, pos[v] + 1)
    return res

模範解答 (Python)

import sys
from math import isqrt
input = sys.stdin.readline

class BIT:
    def __init__(self, n):
        self.n = n
        self.data = [0] * (n + 1)
    def add(self, i, v):
        while i <= self.n:
            self.data[i] += v
            i += i & -i
    def prefix(self, i):
        s = 0
        while i > 0:
            s += self.data[i]
            i -= i & -i
        return s
    def query(self, l, r):
        return self.prefix(r) - self.prefix(l - 1)

def solve():
    N, Q = map(int, input().split())
    parent = [0] * (N + 1)
    edge_w = [0] * (N + 1)
    children = [[] for _ in range(N + 1)]
    for v in range(2, N + 1):
        p, w = map(int, input().split())
        parent[v] = p
        edge_w[v] = w
        children[p].append(v)

    size = [1] * (N + 1)
    depth = [0] * (N + 1)
    heavy = [-1] * (N + 1)
    order = []
    stack = [1]
    visited = [False] * (N + 1)
    while stack:
        u = stack[-1]
        if not visited[u]:
            visited[u] = True
            order.append(u)
            for c in children[u]:
                depth[c] = depth[u] + 1
                stack.append(c)
        else:
            stack.pop()
            p = parent[u]
            size[p] += size[u]
            if heavy[p] == -1 or size[u] > size[heavy[p]]:
                heavy[p] = u

    pos = [0] * (N + 1)
    head = [0] * (N + 1)
    cur_pos = [0]
    stack2 = [(1, 1)]
    while stack2:
        u, h = stack2.pop()
        head[u] = h
        pos[u] = cur_pos[0]
        cur_pos[0] += 1
        for c in children[u]:
            if c != heavy[u]:
                stack2.append((c, c))
        if heavy[u] != -1:
            stack2.append((heavy[u], h))

    bit = BIT(N + 1)
    for v in range(2, N + 1):
        bit.add(pos[v] + 1, isqrt(edge_w[v]))

    def path_sum(u, v):
        res = 0
        while head[u] != head[v]:
            if depth[head[u]] < depth[head[v]]:
                u, v = v, u
            res += bit.query(pos[head[u]] + 1, pos[u] + 1)
            u = parent[head[u]]
        if depth[u] > depth[v]:
            u, v = v, u
        if u != v:
            res += bit.query(pos[u] + 2, pos[v] + 1)
        return res

    out = []
    for _ in range(Q):
        line = list(map(int, input().split()))
        if line[0] == 1:
            _, u, w = line
            bit.add(pos[u] + 1, isqrt(w) - isqrt(edge_w[u]))
            edge_w[u] = w
        else:
            _, u, v = line
            out.append(path_sum(u, v))
    print('\n'.join(map(str, out)))

solve()

Step-by-Step 解説

1辺重みの頂点対応
辺 $(parent[v], v)$ の重みを頂点 $v$ に持たせることで点更新・区間和に帰着。

2HLD 前処理
重い子（サブツリーが最大の子）を優先して DFS し、連続した HLD 位置を割り当てる。

3BIT 初期化
bit.add(pos[v]+1, isqrt(edge_w[v])) で各辺の $\lfloor\sqrt{w}\rfloor$ を登録。

4パスクエリ
チェーンを登りながら BIT 区間和を積算。LCA での辺除外（pos[u]+2 から）に注意。

5差分更新
重み変更時は $\lfloor\sqrt{w_{\text{new}}}\rfloor - \lfloor\sqrt{w_{\text{old}}}\rfloor$ を BIT に add。

計算量

HLD 前処理: $O(N)$
パスクエリ: $O(\log^2 N)$ — チェーン $O(\log N)$ × BIT $O(\log N)$
点更新: $O(\log N)$

よくあるミス

ミス	原因	正しい書き方
LCA の辺を二重カウント	辺クエリで LCA 自身の辺を含める	`pos[u]+2` から始めて LCA 辺を除外
浅い側の頂点に辺を持たせる	辺と頂点の対応ミス	必ず深い側（子）の `pos[v]` に格納
DFS スタックで重辺を後に処理	軽辺の子を後に積むべき	軽辺を先に `stack2.append` してから重辺を積む
更新時に old_sq を再計算忘れ	BIT が古い値のまま	`isqrt(w) - isqrt(edge_w[u])` を add

Day 039-Q2 — HLD + BIT による木パス sqrt 和クエリ

問題

制約