Day 034-Q1 | SegTree上の二分探索 (Walk on Segment Tree)

問題

長さ $N$ の数列 $A_1, A_2, \ldots, A_N$ が与えられる。以下の $Q$ 個のクエリに答えよ。

update l r x $A_l, \ldots, A_r$ の各要素に $x$ を加算する。
query l k 区間 $[l, N]$ で $A_l + \ldots + A_r \geq k$ を満たす最小の $r$ を求めよ。なければ $-1$。

制約

$1 \le N \le 2 \times 10^5$

$1 \le Q \le 2 \times 10^5$

$1 \le A_i \le 10^9$

$-10^9 \le x \le 10^9$

$1 \le k \le 10^{18}$

要素は常に $\ge 0$ を保証

入出力例

入力例 1

5 4
1 2 3 4 5
query 1 6
update 2 4 10
query 2 20
query 3 100

出力例 1

3
3
-1

query 1 6: $1 + 2 + 3 = 6 \ge 6$ → $r=3$
update 2 4 10: 数列 = $[1, 12, 13, 14, 5]$
query 2 20: $12 + 13 = 25 \ge 20$ → $r=3$
query 3 100: $13+14+5 = 32 < 100$ → $-1$

概念図: Walk on Segment Tree

セグメント木の 根から葉まで「左子で目標達成できる？」と問いながら降りる。

黄色矢印 = walk 経路、赤枠 = 答えに到達した葉 (r=3)。各ノードで tree_sum < remaining_k なら -1 を返し、葉なら答えを確定する。

ヒント (段階的開示)

ヒント1: 方向性

区間加算 + 「累積和が閾値以上になる最左位置」クエリ。素朴には $O(N)$、効率化により $O(\log^2 N)$ または $O(\log N)$ で処理する。

ヒント2: アプローチ

遅延伝播セグメント木 で各ノードに区間総和を持ち、クエリは「セグメント木上を歩く (Walk on Segment Tree)」。左の累積和が $k$ に満たなければ右へ、満たすなら左へ降りる二分探索を行う。

ヒント3: 擬似コード

# セグメント木上の歩き方
def walk(node, node_l, node_r, query_l, remaining_k):
    if node_r < query_l or tree_sum[node] < remaining_k:
        return -1
    if node_l == node_r:
        return node_l
    left_result = walk(2*node, node_l, mid, query_l, remaining_k)
    if left_result != -1:
        return left_result
    left_contribution = range_sum(2*node, ..., query_l, mid)
    return walk(2*node+1, mid+1, node_r, query_l, remaining_k - left_contribution)

模範解答 (Python)

import sys
from sys import stdin
input = stdin.readline

def main():
    input_data = sys.stdin.read().split()
    idx = 0

    N, Q = int(input_data[idx]), int(input_data[idx+1])
    idx += 2
    A = [int(input_data[idx+i]) for i in range(N)]
    idx += N

    # 遅延伝播セグメント木 (区間加算 / 区間和)
    size = 1
    while size < N:
        size <<= 1

    tree = [0] * (2 * size)
    lazy = [0] * (2 * size)

    for i in range(N):
        tree[size + i] = A[i]
    for i in range(size - 1, 0, -1):
        tree[i] = tree[2*i] + tree[2*i+1]

    def push_down(node, node_len):
        if lazy[node] != 0:
            half = node_len // 2
            tree[2*node]   += lazy[node] * half
            lazy[2*node]   += lazy[node]
            tree[2*node+1] += lazy[node] * half
            lazy[2*node+1] += lazy[node]
            lazy[node] = 0

    def update(node, node_l, node_r, l, r, val):
        if r < node_l or node_r < l:
            return
        if l <= node_l and node_r <= r:
            tree[node] += val * (node_r - node_l + 1)
            lazy[node] += val
            return
        push_down(node, node_r - node_l + 1)
        mid = (node_l + node_r) // 2
        update(2*node,   node_l,  mid,    l, r, val)
        update(2*node+1, mid+1,   node_r, l, r, val)
        tree[node] = tree[2*node] + tree[2*node+1]

    def range_sum(node, node_l, node_r, l, r):
        if r < node_l or node_r < l:
            return 0
        if l <= node_l and node_r <= r:
            return tree[node]
        push_down(node, node_r - node_l + 1)
        mid = (node_l + node_r) // 2
        return (range_sum(2*node,   node_l, mid,    l, r)
              + range_sum(2*node+1, mid+1,  node_r, l, r))

    def walk(node, node_l, node_r, query_l, k_remain):
        if node_r < query_l:
            return -1
        effective_sum = range_sum(node, node_l, node_r, query_l, node_r)
        if effective_sum < k_remain:
            return -1
        if node_l == node_r:
            return node_l
        push_down(node, node_r - node_l + 1)
        mid = (node_l + node_r) // 2
        left_result = walk(2*node, node_l, mid, query_l, k_remain)
        if left_result != -1:
            return left_result
        left_eff = range_sum(2*node, node_l, mid, query_l, mid)
        return walk(2*node+1, mid+1, node_r, query_l, k_remain - left_eff)

    out = []
    for _ in range(Q):
        op = input_data[idx]; idx += 1
        if op == 'update':
            l, r, x = int(input_data[idx]), int(input_data[idx+1]), int(input_data[idx+2])
            idx += 3
            update(1, 1, size, l, r, x)
        else:
            l, k = int(input_data[idx]), int(input_data[idx+1])
            idx += 2
            out.append(str(walk(1, 1, size, l, k)))

    print('\n'.join(out))

main()

Step-by-Step 解説

1遅延伝播セグメント木の構築
各ノードに区間総和を持つ。lazy[node] = 「このノードの要素全てに加算すべき値」を遅延管理。伝播時は子ノードの和に lazy * 子の長さ を加え、子の lazy に値を伝える。

2Walk on Segment Tree
walk(node, node_l, node_r, query_l, k_remain) はクエリ区間 $[query_l, N]$ で累積和が $k_{remain}$ 以上になる最小位置を返す。

ノードが query_l より左 → -1
有効区間の和が k_remain 未満 → -1
葉 → node_l を返す
それ以外: 左子→ダメなら左子の有効和を引いて右子

3計算量分析
update は $O(\log N)$。walk は各レベルで range_sum を呼ぶため $O(\log^2 N)$。全体 $O((N+Q)\log^2 N)$。

計算量

前処理: $O(N)$
update: $O(\log N)$
walk : $O(\log^2 N)$ （$O(\log N)$ への最適化も可能）
合計: $O((N+Q)\log^2 N)$

よくあるミス

ミス	原因	正しい書き方
push_down を忘れて降りる	遅延値が子に伝わらない	非葉ノードで左右に降りる前に必ず `push_down`
`effective_sum` の計算ミス	`query_l` より左の和を含めてしまう	`range_sum(node, node_l, node_r, query_l, node_r)` で有効範囲のみ
`k_remain` 更新忘れ	左子の和を引かずに右子に渡す	`k_remain - left_eff` を正確に
size > N のパディング非ゼロ	初期化不足	size 以降を 0 で初期化

次のステップ

発展: update を「区間代入」に変更 → Segment Tree Beats と組み合わせた Walk
応用: HLD + Walk による「重み付きグラフでのパスクエリ」

自己評価

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

Day 034-Q1 — SegTree上の二分探索 (Walk on Segment Tree)

問題

制約