公式

E - 家計簿と目標残高 / Household Budget and Target Balance 解説 by admin

gpt-5.5-high

概要

各取引を「入金なら正、出金なら負」の値として扱い、点更新をしながら「累積和が初めて \(K\) 以上になる位置」を答える問題です。

累積和は途中で減ることがあるため、単純な二分探索ではなく、区間の「合計」と「最大接頭辞和」を持つセグメント木を使います。

考察

\(i\) 番目の取引を次の値 \(A_i\) に変換します。

  • R B のとき \(A_i = B\)
  • L B のとき \(A_i = -B\)

すると、\(i\) 番目の取引後の残高は

\[ S_i = A_1 + A_2 + \cdots + A_i \]

です。

タイプ \(2\) の質問では、初めて

\[ S_i \geq K \]

となる \(i\) を求めます。

素朴な方法

質問ごとに \(1\) 番目から順に累積和を計算すると、\(1\) 回の質問に \(O(N)\) かかります。

\(N, Q\) は最大で \(2 \times 10^5\) なので、最悪で

\[ O(NQ) \]

となり間に合いません。

単純な二分探索はできない

累積和 \(S_i\) が単調増加なら二分探索できますが、出金 L があるため累積和は減ることがあります。

例えば、

\[ A = [10, -8, 10] \]

なら累積和は

\[ [10, 2, 12] \]

となり、単調ではありません。

そのため、普通の二分探索では正しく求められません。

重要な気づき

ある区間について、次の \(2\) つを持てばよいです。

  1. 区間全体の合計
  2. 区間内の最大接頭辞和

区間 \([l, r]\) の最大接頭辞和とは、

\[ \max_{l \leq t \leq r}(A_l + A_{l+1} + \cdots + A_t) \]

です。

これが分かれば、「この区間のどこかで残高が \(K\) 以上になるか」を判定できます。

アルゴリズム

セグメント木を使います。

各ノードが表す区間について、以下を管理します。

  • sum:区間全体の合計
  • pref:区間の最大接頭辞和

葉、つまり \(1\) つの要素 \(x\) に対しては、

\[ sum = x \]

\[ pref = x \]

です。

区間のマージ

左区間を \(L\)、右区間を \(R\) とします。

区間全体の合計は単純に、

\[ sum = sum_L + sum_R \]

です。

最大接頭辞和は次のどちらかです。

  1. 左区間の中だけで最大になる
  2. 左区間全体を通って、右区間の途中まで進む

したがって、

\[ pref = \max(pref_L, sum_L + pref_R) \]

となります。

更新クエリ

1 p D b では、\(p\) 番目の値を更新します。

  • R なら \(+b\)
  • L なら \(-b\)

に変更し、セグメント木の葉を更新します。

その後、親方向に sum と pref を再計算します。

これは \(O(\log N)\) です。

質問クエリ

2 K では、初めて累積和が \(K\) 以上になる位置を探します。

まず、全体の最大接頭辞和 seg_pref[1] を見ます。

もし

\[ seg\_pref[1] < K \]

なら、どの取引後にも残高は \(K\) 以上にならないので、答えは \(0\) です。

そうでなければ、セグメント木を上から降りていきます。

現在見ているノードの前までの累積和を acc とします。

左の子について、

\[ acc + pref_{\text{left}} \geq K \]

なら、答えは左の子の区間内にあります。

そうでなければ、左の子の区間内では \(K\) に届かないので、左区間全体を処理したことにして、

\[ acc \leftarrow acc + sum_{\text{left}} \]

として右の子へ進みます。

これを葉に到達するまで繰り返すと、初めて残高が \(K\) 以上になる取引番号が分かります。

計算量

  • 時間計算量: \(O((N+Q)\log N)\)
  • 空間計算量: \(O(N)\)

実装のポイント

  • R は正の値、L は負の値として扱います。

  • セグメント木では各ノードに sum と pref の \(2\) 種類の値を持ちます。

  • 配列サイズを \(2\) のべき乗にそろえるため、実際の取引が存在しない葉があります。

    • その葉の sum は \(0\)
    • pref は十分小さい負の値にして、答えに影響しないようにします。
  • Python では整数が任意精度なのでオーバーフローの心配はありませんが、他の言語では \(64\) bit 整数を使う必要があります。

    ソースコード

import sys

def main():
    input = sys.stdin.buffer.readline

    N, Q = map(int, input().split())

    size = 1
    while size < N:
        size <<= 1

    NEG = -10**30
    seg_sum = [0] * (2 * size)
    seg_pref = [NEG] * (2 * size)

    for i in range(N):
        D, B = input().split()
        v = int(B)
        if D == b'L':
            v = -v
        idx = size + i
        seg_sum[idx] = v
        seg_pref[idx] = v

    for i in range(size - 1, 0, -1):
        l = i << 1
        r = l | 1
        sl = seg_sum[l]
        seg_sum[i] = sl + seg_sum[r]
        a = seg_pref[l]
        b = sl + seg_pref[r]
        seg_pref[i] = a if a >= b else b

    ans = []

    for _ in range(Q):
        q = input().split()

        if q[0] == b'1':
            p = int(q[1])
            D = q[2]
            b = int(q[3])
            if D == b'L':
                b = -b

            idx = size + p - 1
            seg_sum[idx] = b
            seg_pref[idx] = b

            idx >>= 1
            while idx:
                l = idx << 1
                r = l | 1
                sl = seg_sum[l]
                seg_sum[idx] = sl + seg_sum[r]
                a = seg_pref[l]
                c = sl + seg_pref[r]
                seg_pref[idx] = a if a >= c else c
                idx >>= 1

        else:
            K = int(q[1])

            if seg_pref[1] < K:
                ans.append("0")
                continue

            idx = 1
            acc = 0
            while idx < size:
                l = idx << 1
                if acc + seg_pref[l] >= K:
                    idx = l
                else:
                    acc += seg_sum[l]
                    idx = l | 1

            ans.append(str(idx - size + 1))

    sys.stdout.write("\n".join(ans))

if __name__ == "__main__":
    main()

この解説は gpt-5.5-high によって生成されました。