2020浙江工业大学程序设计迎新赛——决赛（重现赛）C-01串(Hard Version)

题目

Description

给定一个只包含 0,10,1 两种字符的字符串 ss ，然后定义一个函数 f(d)f(d) 表示字符串 ss 中距离为 dd 的不同字符对 $(s_i,s_{i+d})$的对数。求 ${\sum }_{d=1}^{k}f(d)$
Input

输入两行。

第一行输入一个正整数 $k$。

第二行输入一个字符串 $s$。

$1\le k<\mid s\mid$
$|s|\le 5\times 10^5$即字符串的长度不超过 $5\times 10^5$
Output
输出一个整数表示 ${\sum }_{d=1}^{k}f(d)$的值。

Sample Input 1

1
0110
Sample Output 1

2
Hint

样例解释：d=1，$(s_0,s_1),(s_2,s_3)$符合条件，合计为 2。

题意

求距离不大于k的所有01对的数量

思路

对于每个0 只需要知道$(i-k,i+k)$范围内1的个数
对于每个1 只需要知道$(i-k,i+k)$范围内0的个数
考虑前缀和预处理0和1的个数，直接$O(n)$扫一遍记录答案，这样会记录两次答案，最后除以2即可。

代码

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

#define moxiaobai false
#define plog if(moxiaobai) cout
typedef long long ll;

const int maxn = 5e5+100;

ll zero[maxn],one[maxn];
int main()
{
	int k;
	string s;
	cin>>k>>s;
	int len = s.size();
	ll z=0,o=0;
	for(int i=0;i<len;i++)
	{
		if(s[i]=='0')
		{
			z++;
			zero[i+1] = z;
			one[i+1] = o;
		}
		if(s[i]=='1')
		{
			o++;
			one[i+1] = o;
			zero[i+1] = z;
		}
	}
	for(int i=0;i<len;i++)
	plog<<zero[i]<<" ";
	ll ans = 0;
	for(int i=0;i<len;i++)
	{
		if(s[i]=='0')
		ans += one[min(len,i+k+1)] - one[max(0,i-k)];
		else 
		ans += zero[min(len,i+k+1)] - zero[max(0,i-k)];
	}
	cout<<ans/2;
}