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;
}