[CF1398C] Good Subarrays

[CF1398C] Good Subarrays

Description

求子串中所有元素的和等于它的长度的子串的数量

Solution

前缀和一下，即 s(r)-s(l)=r-l, s(r)-r=s(l)-l

所以统计出 s(i)-i=k 的数量，然后算答案即可

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

#define int long long

signed main()
{
    ios::sync_with_stdio(false);

    int t;
    cin >> t;

    while (t--)
    {
        int n;
        cin >> n;
        string str;
        cin >> str;
        vector<int> a(n + 2);
        for (int i = 1; i <= n; i++)
            a[i] = str[i - 1] - '0';
        for (int i = 1; i <= n; i++)
            a[i] += a[i - 1];
        map<int, int> mp;
        for (int i = 1; i <= n; i++)
            mp[a[i] - i]++;
        mp[0]++;
        int ans = 0;
        for (auto [x, y] : mp)
            ans += y * (y - 1) / 2;
        cout << ans << endl;
    }
}