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;
}