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 << "
";
            if((pl == i + 1 || pr == i + 1) && k)  continue;
            // cout << 1 << endl;
            // cout << "
";
            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 << "
";
    return 0;
}