Atcoder解いてみた AtCoder Beginner Contest D - Enough Array

AtCoder Beginner Contest D - Enough Array

問題はこちら

回答

尺取り法の練習

K未満の区間を数えて全体から引く
以下、間違ったコード
常に、半壊区間[l,r) l=Kとなる時も、lowerKに1が足されてしまう

#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>

int N;
long long K;
std::vector<long long> A;

int main()
{
	std::cin.tie(0);
	std::ios::sync_with_stdio(false);

	std::cin >> N >> K;
	
	for (int i = 0; i < N; i++) {
		long long a;
		std::cin >> a;
		A.push_back(a);
	}


	int right = 1;
	
	long long lowerK = 0;
	long long sum = A[0];
	for (int left = 0; left < N; left++) {
		// [l,r)を保つ
		if (right == left) {
			right++;
			sum = A[left];
		}


		while (right < N && sum + A[right] < K) {
			sum += A[right];
			right++;
		}

		// sum + A[right]がK以上になってしまったらlowerKは0にしたいが、1が加算されてしまう
		lowerK += (right - left);

		sum -= A[left];
	}

	long long n = static_cast<long long>(N);
	long long ans = (n * (n - 1) / 2 + n) - lowerK;

	std::cout << ans << std::endl;

	return 0;
}

K以下の区間の数を求めるときは、以下のタイプの書き方がいい

#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>

int N;
long long K;
std::vector<long long> A;

int main()
{
	std::cin.tie(0);
	std::ios::sync_with_stdio(false);

	std::cin >> N >> K;

	for (int i = 0; i < N; i++) {
		long long a;
		std::cin >> a;
		A.push_back(a);
	}


	int right = 0;

	long long lowerK = 0;
	long long sum = 0;
	for (int left = 0; left < N; left++) {

		while (right < N && sum + A[right] < K) {
			sum += A[right];
			right++;
		}

		lowerK += (right - left);

		if (left == right) {
			right++;
		}
		else {
			sum -= A[left];
		}
	}

	long long n = static_cast<long long>(N);
	long long ans = (n * (n - 1) / 2 + n) - lowerK;

	std::cout << ans << std::endl;

	return 0;
}

「K未満の区間を数えて全体から引く」よりも、「和が K 以上になる最小の right を探して直接数え上げる」 しゃくとり法にする方が条件分岐がシンプルになり、バグを防ぎやすくなる。

#include <iostream>
#include <vector>

int main() {
    std::cin.tie(0);
    std::ios::sync_with_stdio(false);

    int N;
    long long K;
    std::cin >> N >> K;

    std::vector<long long> A(N);
    for (int i = 0; i < N; i++) {
        std::cin >> A[i];
    }

    long long ans = 0;
    long long sum = 0;
    int right = 0;

    for (int left = 0; left < N; left++) {
        // sum が K 以上になるまで right を進める
        while (right < N && sum < K) {
            sum += A[right];
            right++;
        }

        // sum が K 以上になれば、[left, right) 以降の right~N までの区間はすべて K 以上
        if (sum >= K) {
            ans += (N - right + 1);
        } else {
            // right == N に達しても sum < K の場合はこれ以上 left を進めても K 以上になれないため終了
            break;
        }

        // left を進める準備
        sum -= A[left];
    }

    std::cout << ans << std::endl;

    return 0;
}