AGC052C 题解

### [AGC052C](https://atcoder.jp/contests/AGC052/tasks/agc052_c) 简化题面： 我们总共有 n 个数，每个数的取值是 $[1, p-1]$。 对于 $(p - 1) ^n$ 个数列，称其为好的数列当且仅当存在一种方案使其重排，让其所有前缀都不是 p 的倍数。 求好的数列的个数。 --- 很容易想到对于不要求重排的数列方案使 $(p - 1) \times (p - 2) ^ {n - 1}$ > 除了第一个数随便选，剩下的只有 $(p - 2)$ 种取法。 之后考虑对于答案，既然要求重排肯定会有重复的。 --- 我们考虑放数，当放 x 的时候构成了前缀为 p 的倍数，那么我们考虑放一个 $y, y \ne x$ 之后再放 x，因为 $x \ne y, x +y \ne x$ 所以可以得到一个符合要求的数列。 那么对于一种数列其没法重构符合条件的情况就是 > $1)$ 所有数的权值之和就是 p 的倍数 > $2)$ 当放到最后的时候不能放 x 但是只有 x 可供选择 我们考虑简化这个条件，因为 p 是质数，而我们所有的操作都是在 $\mod p$ 的情况下，所以可以同乘以众数的逆元，因为 $x \times x^{-1} \equiv 1 \pmod p$ 而其他数都不会变成 1。所以我们只需要考虑众数是 1 的情况。 既然众数是 1 我们需要尽量将其向前放，那么我们可以考虑一种方案，就是一开始可以放 $p - 1$ 个 1 之后必须放一个另外的数，我们记为 $b_i$。那么我们可以再放 $p - b_i$ 个 1。 我们在什么情况中是不能满足条件的呢？我们将所有的 $b_i$ 都用上了都不能满足条件。记 $cnt_1$ 表示 1 的个数，$m$ 表示 $n - cnt_1$。满足条件的情况是 $$ \begin{aligned} cnt_1 &\le \sum_{i = 1} ^ m (p - b_i) + (p - 1) \\ cnt_1 &\le (m + 1)p - sum_b - 1 \end{aligned} $$ 发现这个不好求，我们考虑正难则反，求不符合条件的。也就是 $$ \begin{aligned} cnt_1 &> (m +1 )p - sum_b - 1 \\ cnt_1 &\ge (m + 1)p - sum_b \end{aligned} $$ 之后考虑 dp，$dp[i][j]$ 表示已经选了不是众数的 i 个数，$j = \sum_{k = 1}^i (p-b_i)$ 可能发现第二维特别大，但是我们可以发现 $n - i \ge j + p$ > 这里具体就是从之前的式子中推出来的。 $$ \begin{aligned} cnt_1 &\ge \sum_{i = 1} ^ m (p - b_i) + p \\ n - i &\ge j + p \\ j &\le n - i - p \end{aligned} $$ 所以我们可以保证 j 这一维是 $\le n$ 的 > 我们看看之前想求符合条件的，其不行的原因就是 dp 的状态本身就太大了。 可以发现转移是枚举一个之后的 $b_i$ 复杂短 $O(n^3)$ 之后加一个前缀和优化，复杂度 $O(n^2)$ 但是总数怎么算呢？ 我们考虑一下之前的限制，发现还有一种是总和不为 p 的倍数。 那么我们考虑设 $g[i]$ 表示序列长度为 i 其总和不为 p 的方案数。 转移的时候考虑从前 $i - 1$ 个数是 p 的倍数和不是 p 的倍数进行转移。 $$ \begin{aligned} g[i] &= g[i - 1]\times(p - 1) + [(p - 1) ^{i - 1} - g[i - 1]] \times (p - 2) \\ g[i] &= (p - 1) ^ i - g[i - 1] \\ \end{aligned} $$ 这显然是一个递推式子，先算一下前两项 $g[0] = 0, g[1] = p - 1$ 之后考虑展开，可以发现左后一项都是正的，就是 $(p - 1) ^ n - (p - 1)^{n - 1} \dots ...$ 之后最后一项再对于 n 是奇数还是偶数分类讨论就可以了。 > 有能力的同学可以用生成函数 得到： $$ \begin{aligned} s = \frac{(P-1)^N - (P-1)}{P}, n \text{ is odd} \\ s = \frac{(P-1)^N + (P-1)}{P}, n \text{ is even} \end{aligned} $$ 就可以写代码了。 --- $Code$ ```cpp #include <bits/stdc++.h> using namespace std; template <typename T> void r1(T &x) { x = 0; char c(getchar()); int f(1); for(; !isdigit(c); c = getchar()) if(c == '-') f = -1; for(; isdigit(c);c = getchar()) x = (x << 1) + (x << 3) + (c ^ 48); x *= f; } #define int long long #define mod 998244353 const int maxn = 5e3 + 5; const int maxm = maxn << 1; typedef int room[maxn]; int n, p, fac[maxn], inv[maxn]; int C(int a,int b) { // a - b > 0 return fac[a] * inv[b] % mod * inv[a - b] % mod; } int ksm(int x,int mi) { x = (x % mod + mod) % mod; int res(1); while(mi) { if(mi & 1) res = 1ll * res * x % mod; mi >>= 1; x = 1ll * x * x % mod; } return res; } int f[maxn][maxn]; signed main() { int i, j; r1(n), r1(p); fac[0] = 1; for(i = 1; i < maxn; ++ i) fac[i] = 1ll * fac[i - 1] * i % mod; inv[maxn - 1] = ksm(fac[maxn - 1], mod - 2); for(i = maxn - 2; ~ i; -- i) inv[i] = 1ll * inv[i + 1] * (i + 1) % mod; int z = min(p - 2, n); f[0][0] = 1; for(i = 1;i <= n; ++ i) { int sum(0); for(j = 1;j <= n; ++ j) { sum = (sum + f[i - 1][j - 1]) % mod; if(j - z > 0) sum = (sum + mod - f[i - 1][j - z - 1]) % mod; // 减去之前多加的 f[i][j] = sum; } } // puts("ZZZ"); // puts("ZZZ"); int ans = ksm(p - 1, n) , invp = ksm(p, mod - 2); if(n & 1) ans = (ans - 1ll * (ans - (p - 1) + mod) % mod * invp % mod + mod ) % mod; else ans = (ans - 1ll * (ans + (p - 1)) % mod * invp % mod + mod ) % mod; for(i = 0; i <= n; ++ i) { for(j = 0; j <= n; ++ j) { if(((n - i) % p != j % p) && (n - i >= j + p)) { ans = (ans - 1ll * f[i][j] * (p - 1) % mod * C(n, i) % mod + mod) % mod; } } } printf("%lld\n", ans); return 0; } ```