Ncs

题目描述 Problem Description

某巨魔有条长度为 N 的由红色头骨和绿色头骨组成的项链（当然事实上是没有的），这条项链现在已经断开了。 由于这条项链太长了，巨魔想截取其中一段长度非零的项链作为自己的新项链。 这个巨魔很迷信，它认为截取下来的项链段上如果有正好 M 个红色头骨那么这段项链就可以认为是有利于报复社会的。 它想知道有多少种截取方法使得截下来的项链利于报复社会。 两种截取方法被认为不同当且仅当它们至少有一个截取的位置不同。 当然，巨魔也可以保留这条项链，不截取。

输入描述 Input Description

第一行一个数字 M 。

第二行一个 01 字符串， 0 表示绿色头骨， 1 表示红色头骨。 该字符串的长度为 N 。

输出描述 Output Description

输出一行，即截下来的项链有利于报复社会的截取方案数。

输入样例 Sample Input

1

1010

输出样例 Sample Output

6

提示 Hint

截取方案： [1,1],[1,2],[2,3],[2,4],[3,4],[3,3] 。总共 6 种。

数据范围 Data Size

10% 的数据： 1≤N≤100。

30% 的数据： 1≤N≤1000 。

100% 的数据： 1≤N≤1000000。0≤M≤1000000 。

分析 I Think

考虑前缀和，记 Si 表示第 i 位(含)之前的 1 的个数， Ri 表示 Sk=i 的个数。

如果 n=0 ，那么 ans=∑SNi=1Ri∗Ri+M

如果 n≠0 ，那么 ans=∑SNi=1Ri∗(Ri−1)2

代码 Code

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

typedef long long LL;

const int Size = 1000000+47;

char str[Size];
int len;
int sum[Size];
int a[Size];
int n;

int main(){

scanf("%d%s",&n,str+1);

len = strlen(str+1);

a[0] = 1;
for(int i=1;i<=len;++i)
++a[sum[i]=sum[i-1]+str[i]-'0'];

LL ans = 0;
if(n)
for(int i=0;i<=sum[len];++i)
ans += a[i]*a[i+n];
else
for(int i=0;i<=sum[len];++i)
ans += (a[i]*(a[i]-1)>>1);

printf("%lld",ans);

return 0;

}