I am trying to solve 165C using binary search, but it is getting TLE. Can anyone help?
Here's my code: 208389930
Code
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
ll n, k;
ll pre[100005];
string s;
ll bs1(ll lo){
ll l = lo, r = n, mid;
while (l <= r){
mid = (l + r) / 2;
if (pre[mid] — pre[lo-1] <= k){
l = mid + 1;
} else {
r = mid — 1;
}
}
return r;
}
ll bs2(ll lo){
ll l = lo, r = n, mid;
while (l <= r){
mid = (l + r) / 2;
if (pre[mid] — pre[lo-1] >= k){
r = mid — 1;
} else {
l = mid + 1;
}
}
return l;
}
signed main(){
ll i, res;
cin >> k >> s;
n = s.size();
s = " " + s;
pre[0] = 0;
for (i = 1; i <= n; i++){
pre[i] = pre[i-1] + s[i] — '0';
}
res = 0;
for (i = 1; i <= n; i++){
res += bs1(i) — bs2(i) + 1;
}
cout << res;
}