[AtCoder][ABC130D] しゃくとり法の解説

競プロでよく使われるテクニックの１つである「しゃくとり法」について解説する
[ABC130 D - Enough Array] (https://atcoder.jp/contests/abc130/tasks/abc130_d)

問題

配列における連続する部分列の和がK以上になるような組み合わせの数を求めよ

例えば下記の配列AとKを考える。
A = [6 1 2 7], K = 10

単純に考えると２つのポインタ(left, right)を用意してポインタ間の部分列がK以上になるかどうか確認するとO(n^2)の計算量で計算できる。

1. (left=0, right=0) sum = 6
2. (left=0, right=1) sum = 6 + 1 = 7
3. (left=0, right=2) sum = 6 + 1 + 2 = 9
4. (left=0, right=3) sum = 6 + 1 + 2 + 7 = 16 ★ここで初めてKを超えるのでans+1
5. (left=1, right=1) sum = 1
6. (left=1, right=2) sum = 1 + 2 = 3

（省略）

コードで書くと下記のようになる。

int ans = 0;
for (int left=0; left<n; ++left) {
    int sum = 0;
    for (int right=left; right<n; ++right) {
        sum += A[right];
        if (sum >= K) ans++;
    }
}

しゃくとり法

上記問題をよく考えると無駄な部分がある。
部分列の和は単調増加なので、sum(left, right)がK以下であればsum(left+1, right)もK以下となる。上記で5と6の部分はあきらかにK以下となるので省略可能。

よってleftを更新するたびにrightをleftの位置まで戻す必要がなく、rightは前回の位置からの再開で問題ない。

1. (left=0, right=0) sum = 6
2. (left=0, right=1) sum = 6 + 1 = 7
3. (left=0, right=2) sum = 6 + 1 + 2 = 9
4. (left=0, right=3) sum = 6 + 1 + 2 + 7 = 16 ★ここで初めてKを超えるのでans+1
5. (left=1, right=3) sum = 1 + 2 + 7 = 10 ★ans+1
6. (left=2, right=3) sum = 2 + 7 = 9
7. (left=3, right=3) sum = 7

コードで書き直すと下記のようになる。
rightポインタとsumの関係がややこしくなるが、半開区間をイメージすると整理できる。sumを計算した後は必ずrightポインタを進める。

sum = [left, right) = leftからright-1までの合計

int sum = 0;
for (int left=0; left<N; ++left) {
    while (right < N && sum+A[right] < K) {
        sum += A[right];
        right++;
    }
    ans += N - right;
    sum -= A[left];
}

関連問題

[典型90 034 - There are few types of elements] (https://atcoder.jp/contests/typical90/tasks/typical90_ah)