CF EDU 120 D - Shuffle

D - Shuffle

组合数学

记当前枚举的区间从第 i 个 1 到 第 i + k - 1 个 1，记 j = i + k - 1, 那这些 1 可以随意排列的区间为 $[pos[i-1]+1, pos[j+1]-1]$, 设为 $[l,r]$, 这个区间对答案的贡献为 $\binom {r-l+1}k$

但是和上一个区间会有重复，重复的数量为 $\binom {pos[j]-pos[i-1]-1}{k-1}$

#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
#include <cmath>

using namespace std;
typedef long long ll;

const int N = 5e3 + 10;
const int mod = 998244353;

int n, k;
int pos[N], cnt;
ll fac[N], finv[N];

ll qmi(ll a, ll b)
{
	ll ans = 1;
	while(b)
	{
		if (b & 1)
			ans = ans * a % mod;
		b >>= 1;
		a = a * a % mod;
	}
	return ans % mod;
}
void presolve()
{
	fac[0] = finv[0] = 1;
	for (int i = 1; i <= N - 5; i++)
	{
		fac[i] = fac[i-1] * i % mod;
		finv[i] = qmi(fac[i], mod - 2);
	}
}

ll C(int n, int m)
{
	if (m < 0 || n - m < 0)
		return 0;
	return fac[n] * finv[m] % mod * finv[n-m] % mod;
}

int main()
{
	ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
	presolve();
	cin >> n >> k;
	string s;
	cin >> s;
	s = " " + s;
	for (int i = 1; i <= n; i++)
		if (s[i] == '1')
			pos[++cnt] = i;
	if (k == 0 || cnt < k)
	{
		cout << 1 << endl;
		return 0;
	}
	ll ans = 0;
	pos[cnt+1] = n + 1;
	for (int i = 1; i + k - 1 <= cnt; i++)
	{
		int j = i + k - 1;
		int l = pos[i-1] + 1;
		int r = pos[j+1] - 1;
		// cout << l << " " << r << endl;
		ans = (ans + C(r - l + 1, k)) % mod;
		if (i > 1)
			ans = (ans - C(pos[j] - pos[i-1] - 1, k - 1)) % mod;
	}
	cout << (ans + mod) % mod << endl;
	return 0;
}