Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Code

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

long long solve(vector<long long> &a, long long k) {
    int n = a.size();
    vector<long long> pref = a;

    // Convert to prefix sum array.
    for (int i = 1; i < n; i++) {
        pref[i] += pref[i - 1];
    }

    // Prefix sum of positive integers is always sorted.

    // dp[i] is the number of subarrays ending at i with sum at least k.
    vector<int> dp(n, 0);

    for (int i = 0; i < n; i++) {
        int max_unsafe_index_for_start =
            (upper_bound(pref.begin(), pref.begin() + i + 1, pref[i] - k) -
             pref.begin());
        int max_safe_index_for_start = max_unsafe_index_for_start - 1;

        dp[i] = (max_safe_index_for_start - 0 + 1);
        dp[i] += (pref[i] >= k); // Is the subarray [0...i] valid.
    }

    long long ans = accumulate(dp.begin(), dp.end(), 0LL);
    return ans;
}

int main() {
    int n;
    cin >> n;
    long long k;
    cin >> k;

    vector<long long> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }

    auto res = solve(a, k);
    cout << res << endl;

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

long long solve(vector<long long> &a, long long k) {
    int n = a.size();
    vector<long long> pref = a;

    // dp[i] is the number of subarrays starting at i with sum at least k.
    vector<int> dp(n, 0);

    // Use 2 pointers.
    int i = 0, j = 0;
    long long current_window_sum = a[0];
    while (i < n) {
        while (j < n && current_window_sum < k) {
            j++;
            current_window_sum += a[j];
        }

        if (current_window_sum >= k) {
            // The ending index can be [j, j + 1, ... n - 1]
            dp[i] = (n - 1) - (j) + 1;
        } else {
            dp[i] = 0;
        }

        current_window_sum -= a[i];
        i++;
    }

    return accumulate(dp.begin(), dp.end(), 0LL);
}

int main() {
    int n;
    cin >> n;
    long long k;
    cin >> k;

    vector<long long> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }

    auto res = solve(a, k);
    cout << res << endl;

    return 0;
}