Day 135-Q1 | セグメント木上の二分探索（Segment Tree Walk）

問題

長さ $N$ の非負整数列 $A=(A_1,\dots,A_N)$ がある。次の2種類のクエリを $Q$ 個、順番に処理せよ。

1 l r x : 区間 $[l,r]$（両端含む）の各要素に非負整数 $x$ を加算する。
2 l x : 位置 $l$ から右へ累積和を取り、和が初めて $x$ 以上になる最小の位置 $r$（$l\le r\le N$）を出力する。存在しなければ -1。

要素・加算値がすべて非負なので、任意の位置からの累積和は単調非減少。この単調性を使い、セグメント木を根から二分探索的に降りる「Segment Tree Walk」でクエリ2を $O(\log N)$ で処理せよ。

入力形式

N Q
A_1 A_2 ... A_N
query_1
...
query_Q

制約

$1 \le N,Q \le 2\times10^5$

$0 \le A_i \le 10^4$

クエリ1: $1\le l\le r\le N,\ 0\le x\le10^4$

クエリ2: $1\le l\le N,\ 1\le x\le10^{14}$

入出力例

入力例1

出力例1

3
1
-1

概念図

ヒント（段階的開示）

ヒント1（方向性）

区間和と「和が $x$ 以上になる位置」を外側で二分探索すると $O(\log^2 N)$ になる。要素が非負なら、セグメント木の構造そのものをたどることで $O(\log N)$ に落とせる。

ヒント2（アプローチ）

「残り必要量 need」を管理しながら根から降下し、左の子の和が need 以上なら左へ、そうでなければ need から左の和を引いて右へ進む。遅延伝播がある場合は降りる前に必ず push down する。

ヒント3（誘導）

区間 $[l,N]$ を通常の区間分解で $O(\log N)$ 個のノードに分け、左から見ていき、和が need を超えたノードの内部だけを降下する2段階探索にする。

def descend(node, node_l, node_r, need):
    if node_r - node_l == 1:
        need[0] -= tree[node]
        return node_l
    push_down(node)
    mid = (node_l + node_r) // 2
    if tree[2*node] >= need[0]:
        return descend(2*node, node_l, mid, need)
    need[0] -= tree[2*node]
    return descend(2*node+1, mid, node_r, need)

模範解答 (Python)

import sys


def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    Q = int(input_data[idx]); idx += 1
    A = [int(input_data[idx + i]) for i in range(N)]
    idx += N

    size = 1
    while size < N:
        size *= 2

    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 apply_add(node, node_len, val):
        tree[node] += val * node_len
        if node < size:
            lazy[node] += val

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

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

    def descend(node, node_l, node_r, need, node_len):
        if node_r - node_l == 1:
            need[0] -= tree[node]
            return node_l
        push_down(node, node_len)
        mid = (node_l + node_r) // 2
        half = node_len // 2
        if tree[2 * node] >= need[0]:
            return descend(2 * node, node_l, mid, need, half)
        need[0] -= tree[2 * node]
        return descend(2 * node + 1, mid, node_r, need, half)

    def query_first(node, node_l, node_r, l, r, need, node_len):
        if r <= node_l or node_r <= l or tree[node] < need[0]:
            if l <= node_l and node_r <= r:
                need[0] -= tree[node]
            return -1
        if l <= node_l and node_r <= r and node_r - node_l == 1:
            return node_l
        push_down(node, node_len)
        mid = (node_l + node_r) // 2
        half = node_len // 2
        if l <= node_l and node_r <= r:
            return descend(node, node_l, node_r, need, node_len)
        res = query_first(2 * node, node_l, mid, l, r, need, half)
        if res != -1:
            return res
        return query_first(2 * node + 1, mid, node_r, l, r, need, half)

    out = []
    for _ in range(Q):
        t = int(input_data[idx]); idx += 1
        if t == 1:
            l = int(input_data[idx]) - 1; idx += 1
            r = int(input_data[idx]); idx += 1
            x = int(input_data[idx]); idx += 1
            range_add(1, 0, size, l, r, x, size)
        else:
            l = int(input_data[idx]) - 1; idx += 1
            x = int(input_data[idx]); idx += 1
            need = [x]
            res = query_first(1, 0, size, l, N, need, size)
            out.append(str(res + 1) if res != -1 else "-1")

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


main()

Step-by-Step 解説

1なぜ二重二分探索では遅いか
外側で $r$ を二分探索、内側で区間和クエリ($O(\log N)$)を呼ぶと合計 $O(\log^2 N)$。Segment Tree Walk なら木の降下1回で $O(\log N)$。

2「範囲の分解」と「降下」を分離
query_firstは通常の区間分解を行い、範囲全体を覆うノードに到達した時点でdescendに切り替えてリーフまで一直線に降りる。

3遅延伝播との整合性
子を参照する前に必ずpush_downする。忘れると未反映の古い和で判定してしまう。

4非負性が単調性の根拠
tree[2*node] >= need[0]で左右を決められるのは要素が非負で区間和が単調非減少だから。

5見つからないケース
範囲全体の和が need 未満なら -1。部分的にしか範囲に入らないノードでは need を減算しない点に注意。

よくあるミス

ミス	原因	正しい書き方
push_downを降下前に呼び忘れる	遅延加算が子に反映されないまま参照	子を見る直前に必ずpush_downする
need[0]を整数で渡す	Pythonの整数は不変なので再帰内更新が反映されない	リストなどミュータブルなオブジェクトで包む
範囲外ノードでもneedを減算	分岐条件を誤る	完全に範囲外なら何もせず-1を返す
木サイズを2べきに拡張し忘れる	size=Nのまま実装	N以上最小の2べきにし余った葉は0で埋める
-1判定漏れ	範囲和 < x でも位置を返してしまう	探索前に範囲和を確認するか降下条件を厳密にする

次のステップ

発展: 負の値を許す場合、区間内の最大値プレフィックスを別途管理すれば木の降下で処理できる場合がある。境界条件を検討せよ。
発展2: クエリ2を「$l$から左方向」に一般化し、両方向対応のSegment Tree Walkを実装せよ。
次回予告: 次回のテーマは今後のセッションで選定

自己評価

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

Day 135-Q1 — セグメント木上の二分探索（Segment Tree Walk）

問題