Official

E - 家計簿と目標残高 / Household Budget and Target Balance Editorial by kyopro_friends Editorial for printing


この問題はセグメントツリーを用いて解くことができます。

\(i\) 回目の取引による残高の変化量を \(A_i\) とすると、問題は次の通りです。

問題:
長さ \(N\) の数列 \((A_1,\dots,A_N)\) が与えられる。以下の \(2\) 種類のクエリを計 \(Q\) 個処理せよ。

  • \(A_i\) を \(x\) に変更する
  • \(A_1+A_2+\ldots+A_i \geq K\) となる最小の \(i\) を求める

2番目のクエリに答えるために、答えを二分探索したくなります。つまり、prefix sum の最大値 \(\max(0, A_1, A_1+A_2,\ldots,A_1+\ldots+A_i)\) を管理したくなります。

区間和に関するクエリと1点更新に関するクエリを同時に処理することができるデータ構図の筆頭はセグメントツリーなので、これで処理できるかを考えます。

ここで例えば

\(\max(0, A_1, A_1+A_2, A_1+A_2+A_3, A_1+A_2+A_3+A_4) \\ =\max(\max(0, A_1, A_1+A_2),\max( A_1+A_2+A_3, A_1+A_2+A_3+A_4)) \\ =\max(\max(0,A_1,A_1+A_2), (A_1+A_2)+\max(0, A_3,A_3+A_4))\)

となることから、2つの区間をマージして prefix sumの最大値を求めるには、それぞれの区間のprefix sumの最大値と区間和が分かれば十分です。よって、その 2 つの情報をセグメントツリーに乗ることでこの問題を解くことができます。

計算量は \(O(N+Q\log N)\) です。

実装例 (C++)

#include<bits/stdc++.h>
#include<atcoder/segtree>
using namespace std;

using S = array<long long, 2>;
S op(S x, S y){ return {x[0]+y[0], max(x[1], x[0]+y[1])}; }
S e(){ return {0, 0}; }

int main(){
  int n, q;
  cin >> n >> q;
  vector<S> a;
  for(int i=0; i<n; i++){
    char d;
    int b;
    cin >> d >> b;
    if(d == 'R'){
      a.push_back({b, b});
    }else{
      a.push_back({-b, 0});
    }
  }

  atcoder::segtree<S, op, e> seg(a);
  for(int i=0; i<q; i++){
    int t;
    cin >> t;
    if(t==1){
      char d;
      int p, b;
      cin >> p >> d >> b;
      if(d == 'R'){
        seg.set(p-1, {b, b});
      }else{
        seg.set(p-1, {-b, 0});
      }
    }else{
      long long k;
      cin >> k;
      auto f=[&](S x){return x[1] < k;};
      int r = seg.max_right(0, f);
      if(r == n){
        cout << 0 << endl;
      }else{
        cout << r+1 << endl;
      }
    }
  }
}

実装例 (Python)

from atcoder.segtree import SegTree

N, Q = map(int, input().split())
A = []
for i in range(N):
  D, B = input().split()
  B = int(B)
  if D == 'R':
    A.append((B, B))
  else:
    A.append((-B, 0))

def op(x, y):
  return (x[0]+y[0], max(x[1], x[0]+y[1]))

seg = SegTree(op, (0,0), A)
for _ in range(Q):
  t, *q = input().split()
  if t == '1':
    P, D, B = q
    P = int(P)
    B = int(B)
    if D == 'R':
      seg.set(P-1, (B, B))
    else:
      seg.set(P-1, (-B, 0))
  else:
    K = int(q[0])
    R = seg.max_right(0, lambda x: x[1] < K)
    if R == N:
      print(0)
    else:
      print(R+1)

posted:
last update: