けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

鉄則本 B13 - Supermarket 2 (2Q, ★4)

しゃくとり法の練習問題!

問題概要

正の整数からなる長さ  N の数列  A_{1}, A_{2}, \dots, A_{N} が与えられる。

この数列の連続する区間であって、総和が  K 以下であるものの個数を求めよ。

制約

  •  1 \le N \le 10^{5}
  •  1 \le K, A_{i} \le 10^{9}

解法 (1):しゃくとり法

editorial に詳しく書かれている。

github.com

コード

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

int main() {
    long long N, K;
    cin >> N >> K;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];
    
    long long res = 0, sum = 0, r = 0;
    for (int l = 0; l < N; ++l) {
        while (r < N && sum + A[r] <= K) {
            sum += A[r++];
        }
        res += r - l;

        // 左端を削る
        sum -= A[l];
    }
    cout << res << endl;
}

解法 (2):累積和上の lower_bound()

もう 1 つの方法として、累積和を取る方法がある。数列  A の累積和を  S とする。つまり、数列  A を 0-indexed で表現したときに、

  •  S_{0} = 0
  •  S_{1} = A_{0}
  •  S_{2} = A_{0} + A_{1}
  •  S_{3} = A_{0} + A_{1} + A_{2}
  •  \dots
  •  S_{N} = A_{0} + A_{1} + A_{1} + \dots + A_{N-1}

とする。このとき、問題は次のように言い換えられる。


 S_{j} - S_{i} \le K となるような組  (i, j)(0 \le i \lt j \le N) の個数を求めよ。


これは lower_bound() などを用いて解決できる有名問題である。