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

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

AtCoder ABC 130 D - Enough Array (緑色, 400 点)

AtCoder AtCoder400点 ABC-D しゃくとり法 二分探索 累積和 区間 数列 解空間:O(N^2)通りの選択肢 二分探索:lower_bound 緑色diff 数え上げ問題 制約条件:総和<=K 連続部分列を扱う問題 解空間:O(N^2)個の区間

超典型的なしゃくとり法そのものだった！！！

問題へのリンク

問題概要

N 個の整数 A_1, A_2, \dots, A_N があたえられる。

数列の連続した区間であって、その総和が K 以上となっているものを数え上げよ。

制約

  • 1 \le N \le 10^{5}

考えたこと

しゃくとり法については以前に記事を書いた。

qiita.com

このしゃくとり記事の 2 番目の例題は

連続した区間であって、その総和が K 以上であるもののうち、区間の長さの最小値を求めよ

という問題であった。今回の問題はこれとほとんど同じ問題である。すなわち「区間の長さの最小値」の代わりに「区間の個数」を求める問題である。

そもそもしゃくとり法を用いることで

区間の左端 left を固定したとき、区間 [left, right) の総和が K 以上となるような最小の right

が求められるのである。

  • 「条件を満たす区間の長さの最小値」だったら、この right - left が答えの候補で、これを全 left について調べれば良い。

  • 「条件を満たす区間の個数」だったら各左端 left に対して、右端が right, right + 1, ..., N (N - right + 1 個) が条件を満たすので、これを各 left について総和をとればよい。

コード

#include <iostream>
#include <vector>
using namespace std;

int N;
long long K;
vector<long long> a;
 
int main() {
    cin >> N >> K;
    a.resize(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
 
    long long res = 0;
        
    /* 区間の左端 left で場合分け */
    int right = 0;
    long long sum = 0;
    for (int left = 0; left < N; ++left) {
        /* [left, right) の総和が K 以上となる最小の right を求める */
        while (right < N && sum < K) {
            sum += a[right];
            ++right;
        }
        if (sum < K) break; // これ以上 left を進めてもダメ

        // 集計
        res += (long long)(N - right + 1);

        if (right == left) ++right;
        else sum -= a[left];
    }
 
    cout << res << endl;
}