Another Problem on Strings
Another Problem on Strings
思路
题目的这一句话提示着我要用binary_search,没想到还真是,
这题考察的就是两个知识点,前缀和以及二分查找,二分查找用STL解决了,前缀和也简单,接下来就是两者组合了。
我们假定当前位置的前缀和是 \(sum[i]\),我们想办法从前面求以\(j\)为起点\(i\)为终点的序列,一定有 \(sum[i] - sum[j] = k\),接下来的事情就简单了,只需要求得\(j\)的区间就行了
我们分别对\(sum[i] - k和sum[i] - k + 1进行lower\_bound\)求得区间 \((l, r]\)正是\(j\)对应的区间。然后就\(ans += r - l\)
然后自信的交了一发,没想到\(wa on test 4\),没错,只过了样例。
这道题目还有一个坑点就是k = 0的时候,上面的方法不能用。所以我就单独分开了一个k = 0的情况讨论。
这个情况比较简单,取 \(i\) 的时候一定要注意这一位是0,然后就是用\(lower \_bound(sum[i])求得l\),其 \(j\) 的取值区间就是\((l, i]\),答案就是\(ans += i - l\)
代码
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 1e6 + 10;
char str[N];
int sum[N], k, n;
int main() {
// freopen("in.txt", "r", stdin);
ios::sync_with_stdio(false);
cin >> k;
cin >> (str + 1);
int n = strlen(str + 1);
for(int i = 1; i <= n; i++)
sum[i] = sum[i - 1] + str[i] - '0';
ll ans = 0;
if(k) {
for(int i = 1; i <= n; i++) {
int l = sum[i] - k, r = sum[i] - k + 1;
// cout << l << " " << r << endl;
int pl = lower_bound(sum, sum + 1 + i, l) - sum;
int pr = lower_bound(sum, sum + 1 + i, r) - sum;
// cout << pl << " " << pr << endl;
// cout << "\n";
if((pl == i + 1 || pr == i + 1) && k) continue;
// cout << 1 << endl;
// cout << "\n";
ans += pr - pl;
}
}
else {
for(int i = 1; i <= n; i++) {
if(str[i] == '1') continue;
int l = lower_bound(sum, sum + 1 + i, sum[i]) - sum;
ans += i - l;
}
}
cout << ans << "\n";
return 0;
}