[CF1188E] Problem from Red Panda 题解

考虑每个位置的操作次数 c_i，不难发现，i 气球最后的颜色个数 d_i 是 a_i + c_ik-\sum c_i，如果存在 \forall c_i > 0，那么我们总是可以把所有气球少操作一次，这样上式不变，不影响最后的序列，下文所有的操作序列都假设 \min c_i = 0。

考虑原序列和操作序列的关系，断言：操作序列（假设后）和原序列一一对应。

证明： 形式化地说，假设有操作序列 A, B, |A| = |B|，最后得到的气球序列为 ta, tb，欲证 A = B \Longleftrightarrow ta = tb。

充分性是显然的，论证必要性：考察气球序列中的位置 i，有 ta_i = tb_i，有 a_i + A_ik-|A| = a_i + B_ik - |B|，因为 |A| = |B|，所以 A_i = B_i，证毕。

得到这个之后，我们只需要计数操作序列即可，并且不用管算重的问题。

发现操作总数 s 总是 \le \max a_i，这是因为如果 s> \max a_i，再根据 \min c_i = 0 的假设，总是存在一个 i，使得 d_i < 0，这非法，所以可以考虑枚举 s。

考虑操作序列合法的充要条件，首先显然有 \forall d_i = a_i + c_ik-s \ge 0，得到 c_i\ge \lceil \dfrac{s - a_i}{k} \rceil，设这个下界是 bound_s(i)，此时我们有 s =\sum c_i \ge \sum bound_s(i)，为了保证任意时刻都存在合法的操作，另一个条件是 \forall t\le s, t\ge\sum bound_t(i)。

证明： 必要性显然。

关于充分性，考虑构造一组合法操作序列，不难发现，随着 t 的增大，bound_t(i) 一定单调不降，对于一个时刻，一定有一部分操作是必须要放到一些位置上，为了满足下界的需求，剩余有一些操作是可以随便放的，所以每次 bound_t(i) 的增加时，把本来可以随便放的操作放到增加的 bound_t(i) 的位置即可，因为有 t\ge\sum bound_t(i)，所以总是能够满足 bound_t(i) 的需求。

根据不等式 c_i\ge \lceil \dfrac{s - a_i}{k} \rceil，如果 a_i \ge s，则 i 位置不需要操作也行，它的 bound_i = 0，而 a_i < s 的部分就有操作的需求，不妨假设 a 不降，随着 s 的增加，有操作需求的位置的分界线 r 会单调地向右移动，用双指针维护这个分界线，同样的，我们也可以维护 w = \sum bound_t(i)，因为每次变化的只有 s - 1\equiv a_i\pmod k 的那些 i，它们的下界会增加 1，用桶维护 w。

剩下的部分是一些简单的组合计数，需要满足：前 r 个位置的操作次数 \ge bound_s(i), i\le r，后 k - r 个位置中至少有一个位置操作次数是 0（需要满足一开始的假设才不会算重），现在有 s 个操作次数，问有几种分配操作次数的方案。

容斥后 k - r 至少有一个 0，然后就是经典小球盒子插板法解决的问题了。

这里的答案就等于：

需要注意的是，如果一个时刻 s < w，直接结束枚举，这是为了满足充要条件 2。

代码实现

#include <iostream>
#include <cstring>
#include <algorithm>
#include <queue>
//#define int long long
#define x first
#define y second
using namespace std;
typedef pair<int, int> PII;
const int N = 2e6 + 10, mod = 998244353;

int k, a[N], fac[N], ifac[N], cnt[N], m;
int qmi(int a, int b) {
    int res = 1;
    while(b) {
        if(b & 1) res = 1ll * res * a % mod;
        a = 1ll * a * a % mod, b >>= 1;
    }
    return res;
}
int C(int n, int m) {
    if(n < m) return 0;
    return 1ll * fac[n] * ifac[m] % mod * ifac[n - m] % mod;
}

signed main() {
    ios::sync_with_stdio(0), cin.tie(0);
    cin >> k;
    for(int i = 1; i <= k; i ++) cin >> a[i], m = max(m, a[i]);
    fac[0] = ifac[0] = 1;
    for(int i = 1; i < N; i ++) fac[i] = 1ll * fac[i - 1] * i % mod;
    ifac[N - 1] = qmi(fac[N - 1], mod - 2);
    for(int i = N - 2; i; i --) ifac[i] = 1ll * ifac[i + 1] * (i + 1) % mod;
    int r = 0, w = 0;
    long long ans = 0;
    sort(a + 1, a + k + 1);
    for(int s = 0, j = 1; s <= m; s ++) {
        while(j <= k && a[j] < s) cnt[a[j] % k] ++, j ++;
        w += cnt[(s - 1 + k) % k];
        if(w > s) break;
        r = j - 1;
        (ans += C(s - w + k - 1, k - 1) - C(s - w + r - 1, k - 1)) %= mod;
    }
    cout << (ans + mod) % mod << '\n';

    return 0;
}