[CF1188E] Problem from Red Panda 题解

[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，然后就是经典小球盒子插板法解决的问题了。

这里的答案就等于：

\[\binom{n - w + k - 1}{k - 1} - \binom{n - w + r - 1}{k - 1} \]

需要注意的是，如果一个时刻 \(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;
}