[组合数学][CF1188E]Problem from Red Panda

题目链接

题目传送门

简要题意

有一个长度为 (k) 的数组 (a)，每次可以选择一个 (1le ile k)，让 (a_i) 加上 (k-1)，并对于所有的 (j e i) 让 (a_j) 减掉 (1)，任何时候必须保证 (a) 数组非负。

求通过任意多次（可以为 (0) 次）操作，能达到的不同的 (a) 数组方案数膜 (998,244,353) 后的结果。

数据范围：(2le kle 10^5)，(a_ige 0)，(sum_{i=1}^ka_ile10^6)，但实际上存在复杂度与 (sum_{i=1}^ka_i) 无关的做法。

Solution

设 (x_i) 表示 (a_i) 被选中的次数，考虑一个 ({x_i}) 合法的条件。记 (s=sum_{i=1}^kx_i)。

我们不妨把操作看作数组 (a) 整体减 (1) 之后 (a_i+=k)。

显然我们必须保证最后的 (a) 数组非负，故 (a_i-s+kx_ige 0)，也就是 (x_igelceilfrac{max(s-a_i,0)}k ceil)。

在这个条件下，判断是否对于 (0le t<s) 满足 (t) 轮操作之后 (a) 数组非负，只需将所有 (a_i) 减掉 (t)，然后尝试用 (t) 个 (k) 来填充为负的 (a_i) 值，判断是否能够填充成功即可，即 (sum_{i=1}^klceilfrac{max(t-a_i,0)}k ceille t)。显然在 (x_igelceilfrac{max(s-a_i,0)}k ceil) 的限制下，第 (i) 个数被操作的次数不会超过 (x_i)。

于是一个 ({x_i}) 合法的条件为：

（1）(x_igelceilfrac{max(s-a_i,0)}k ceil)

（2）对于所有 (0le tle s) 都有 (sum_{i=1}^klceilfrac{max(t-a_i,0)}k ceille t)。

在从小到大枚举 (s) 的过程中，(sum_{i=1}^klceilfrac{max(s-a_i,0)}k ceil) 容易求出，将 (a) 排序后用指针维护 (a_i<s) 的部分，对 (a_imod k) 用个桶维护每种值的出现次数即可。

回到问题，我们不能直接对 ({x_i}) 计数，因为不同的 (x) 数组可能对应同一个 (a) 数组。

首先我们发现，如果所有的 (x_i) 都相等，则这样的操作对原数组没有影响。

也就是说，如果所有的 (x_i) 都不为 (0)，则把所有 (x_i) 都减掉 (1) 之后会得到一个等价的方案。

同样地如果将一部分 (x_i)（个数在 ([1,k-1]) 之间）减掉 (1)，则得到的方案一定不等价。

故可以转化成对 ({x_i}) 数组计数，但 (x) 数组必须满足至少有一个 (0)。易得这时有 (0le slemax a_i)。

从小到大枚举 (s)，遇到 (w=sum_{i=1}^klceilfrac{max(s-a_i,0)}k ceil>s) 的情况立刻 break 掉。

问题转化成 (k) 个变量，其中前 (r)（(a_i<s) 的 (i) 个数）个变量有一个取值下界 (down_i)，满足 (down_ige 1) 且 (w=sum_{i=1}^rdown_i)，求为这 (k) 个变量取值，使得至少有一个 (0)，并且所有变量的和为 (s) 的方案数。

先去掉下界 (down)，转成所有变量的和为 (s-w)，并且后 (k-r) 个变量至少有一个 (0)。

考虑容斥，用任意方案减掉没有 (0) 的方案，由插板法得方案数：

[inom{s-w+k-1}{k-1}-inom{s-w+r-1}{k-1} ]

总复杂度 (O(klog k+max a_i))。

Code

#include <bits/stdc++.h>

template <class T>
inline void read(T &res)
{
	res = 0; bool bo = 0; char c;
	while (((c = getchar()) < '0' || c > '9') && c != '-');
	if (c == '-') bo = 1; else res = c - 48;
	while ((c = getchar()) >= '0' && c <= '9')
		res = (res << 3) + (res << 1) + (c - 48);
	if (bo) res = ~res + 1;
}

const int N = 2e6 + 5, djq = 998244353;

int k, n, a[N], cnt[N], fac[N], inv[N], ans;

int C(int n, int m) {return 1ll * fac[n] * inv[m] % djq * inv[n - m] % djq;}

int main()
{
	fac[0] = inv[0] = inv[1] = 1;
	for (int i = 1; i < N; i++) fac[i] = 1ll * fac[i - 1] * i % djq;
	for (int i = 2; i < N; i++) inv[i] = 1ll * (djq - djq / i) * inv[djq % i] % djq;
	for (int i = 2; i < N; i++) inv[i] = 1ll * inv[i] * inv[i - 1] % djq;
	read(k); int cur = 0;
	for (int i = 1; i <= k; i++) read(a[i]), n += a[i];
	std::sort(a + 1, a + k + 1);
	for (int i = 0, j = 1; i <= a[k]; i++)
	{
		while (a[j] < i) cnt[a[j++] % k]++;
		cur += cnt[(i - 1 + k) % k];
		if (cur > i) return std::cout << ans << std::endl, 0;
		ans = (ans + C(i - cur + k - 1, k - 1)) % djq;
		if (i - cur + j - 2 >= k - 1)
			ans = (ans - C(i - cur + j - 2, k - 1) + djq) % djq;
	}
	return std::cout << ans << std::endl, 0;
}