codeforce 165C Another Problem on Strings

原题链接

题意

给定一个只含0或者1的字符串，给定整数K(0<=K<=1e6)，求字符串中1的数目恰好为K的子串数量。

题解

形如100……001这样的子串符合题意，在两端加上0所得到的子串也符合题意，很明显是乘法原理

要求子串中1的个数恰好为K个，所以枚举子串起点，终点也就随之固定

预处理出1的位置数组id，用pre[i]记录第i个1和第i-1个1之间的0的个数，然后扫一遍OK。

有个坑，k==0需要额外处理

#include<bits/stdc++.h>

using namespace std;
typedef long long LL;
const int maxn=1e6;

char s[maxn+5];
LL pre[maxn+5];
int id[maxn+5];

int main(void)
{
	int k;
	scanf("%d",&k);
	
	scanf("%s",s);
	
	int len=strlen(s);
	int cnt=0;
	for (int i=0;i<=len-1;++i)
	{
		if (s[i]=='1') id[++cnt]=i;
	}
	
	id[0]=-1;
	id[cnt+1]=len;
	for (int i=1;i<=cnt+1;++i)
	{
		pre[i]=id[i]-id[i-1]-1;
	}
	
	if (cnt<k)
	{
		printf("0\n");
		return 0;
	}
	else if (k>=1)
	{
		LL ans=0;
		for (int s=1;s<=cnt-k+1;++s)
		{
			int t=s+k-1;
			ans+=(pre[s]+1)*(pre[t+1]+1);
		}
		printf("%I64d\n",ans);
	}
	else if (k==0)
	{
		LL ans=0;
		for (int i=1;i<=cnt+1;++i)
		{
			LL temp=(pre[i]*(pre[i]+1))/2;
			ans+=temp;
		}
		printf("%I64d\n",ans);
	}
}