Good Subarrays 前缀和+map优化 题解

题目描述

Problem Description

You are given an array a 1 , a 2 , … , a n a_1,a_2,…,a_n a1​,a2​,…,an​ consisting of integers from 0 0 0 to 9 9 9. A subarray$ a_l,a_{l+1},a_{l+2},…,a_{r−1},ar$ is good if the sum of elements of this subarray is equal to the length of this subarray ( ∑ i = l r a i = r − l + 1 \sum_{i = l} ^ r a_i = r - l + 1 ∑i=lr​ai​=r−l+1).

For example, if a = [ 1 , 2 , 0 ] a=[1,2,0] a=[1,2,0], then there are 3 good subarrays: a 1 … 1 = [ 1 ] , a 2 … 3 = [ 2 , 0 ] a 1 … 1 = [ 1 ] , a 2 … 3 = [ 2 , 0 ]   a n d   a 1 … 3 = [ 1 , 2 , 0 ] a 1 … 3 = [ 1 , 2 , 0 ] . a_{1…1}=[1],a_{2…3}=[2,0]a_{1…1}=[1],a2…3=[2,0] \ and\ a_{1…3} =[1,2,0]a_{1…3}=[1,2,0]. a1…1​=[1],a2…3​=[2,0]a1…1​=[1],a2…3=[2,0] and a1…3​=[1,2,0]a1…3​=[1,2,0].

Calculate the number of good subarrays of the array a.

Input

The first line contains one integer t ( 1 ≤ t ≤ 1000 1≤t≤1000 1≤t≤1000) — the number of test cases.

The first line of each test case contains one integer n ( 1 ≤ n ≤ 1 0 5 1≤n≤10^5 1≤n≤105) — the length of the array a.

The second line of each test case contains a string consisting of n decimal digits, where the i − t h i-th i−th digit is equal to the value of a i a_i ai​.

It is guaranteed that the sum of n over all test cases does not exceed 1 0 5 10^5 105.

Output

For each test case print one integer — the number of good subarrays of the array a.

Sample Input

3
3
120
5
11011
6
600005

Sample Output

3
6
1

题目分析

题目大意：给定一个已知长度的序列，求满足区间和等于区间长度的子序列的个数。

因为要求出区间和，且数据不是很强，因此先考虑前缀和优化，求出给定序列的前缀和。

不妨设 s u m [ N ] sum[N] sum[N]表示序列 a a a的前缀和，那么区间 [ l , r ] [l, r] [l,r]内的和可表示为 s u m [ r ] − s u m [ l − 1 ] sum[r] - sum[l - 1] sum[r]−sum[l−1]，那么问题就被转换成了在前缀和 s u m sum sum数组中寻找存在多少对 l , r l, r l,r，使得 s u m [ r ] − s u m [ l − 1 ] = r − l + 1 sum[r] - sum[l - 1] = r - l + 1 sum[r]−sum[l−1]=r−l+1

直接写两层循环，大概模了几组强数据，不出意料的TLE了。考虑对查找的过程进行优化。

s u m [ r ] − s u m [ l − 1 ] = r − l + 1 sum[r] - sum[l - 1] = r - l + 1 sum[r]−sum[l−1]=r−l+1可以变形为 s u m [ r ] − s u m [ l − 1 ] − ( r − l + 1 ) = 0 sum[r] - sum[l - 1] - (r - l + 1) = 0 sum[r]−sum[l−1]−(r−l+1)=0

我们可以将 s u m sum sum数组的每个元素-1，那么式子就转化为了 s u m [ r ] − s u m [ l − 1 ] = 0 sum[r] - sum[l - 1] = 0 sum[r]−sum[l−1]=0

这里再次用到STL中的 m a p map map容器，对空间进行优化。

AC Code

#include 
#define ll long long
using namespace std;
const int N = 1e5 + 10;
char s[N];
int a[N], sum[N];

int main(){
    ios_base::sync_with_stdio(0);
    int t; cin >> t;
    while (t--){
        int n;
        cin >> n >> s + 1;
        for (int i = 1; i