luoguP4491 [HAOI2018]染色 广义容斥原理 + FFT

题目中求恰好\(k\)，那么考虑求至少\(k\)

记\(g(k)\)表示至少\(k\)中颜色出现了恰好\(S\)次

那么，\[g(k) = \binom{M}{k} \frac{N!}{(S!)^k (N-Sk)!} * (M-k)^{N-Sk}\]

根据广义容斥原理，记\(f(i)\)表示恰好\(k\)种颜色出现了恰好\(k\)次

那么，\[f(i) = \sum \limits_{k = i}^M (-1)^{k - i} \binom{k}{i} g(k)\]

化成卷积式

\[f(i) * i! = \sum \limits_{k = i}^M \frac{(-1)^{k - i}}{(k - i)!} k! g(k)\]

令\(F_i = \frac{(-1)^{i}}{i!}\)，\(G_i = i! g(i)\)

记\(H_i\)表示\(f(i) * i\)，那么

\[H_i = \sum \limits_{j = i}^M F(k - i) * G(k)\]

反转下标，有

\[H_{n - i}' = \sum \limits_{i = 0}^{n - i} F(k) * G'(n - i - k)\]

\(NTT\)即可，复杂度\(O(n \log n)\)

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;

#define ri register int
#define rep(io, st, ed) for(ri io = st; io <= ed; io ++)

#define gc getchar
inline int read() {
    int p = 0, w = 1; char c = gc();
    while(c > '9' || c < '0') { if(c == '-') w = -1; c = gc(); }
    while(c >= '0' && c <= '9') p = p * 10 + c - '0', c = gc();
    return p * w;
}

const int sid = 3e5 + 5;
const int cid = 1e7 + 5;
const int mod = 1004535809;

inline int mul(int a, int b) { return 1ll * a * b % mod; }
inline int fp(int a, int k) {
    int ret = 1;
    for( ; k; k >>= 1, a = mul(a, a))
        if(k & 1) ret = mul(ret, a);
    return ret;
}

int N, M, S, n, lg;
int fac[cid], inv[cid];
int rev[sid], f[sid], g[sid], w[sid], W[sid];

inline int C(int n, int m) {
    if(n < m) return 0;
    return mul(fac[n], mul(inv[m], inv[n - m]));
}

inline void NTT(int *a) {
    for(ri i = 0; i < n; i ++)
        if(i < rev[i]) swap(a[i], a[rev[i]]);
    for(ri i = 1; i < n; i <<= 1)
    for(ri j = 0, kj = n / (i << 1); j < n; j += (i << 1))
    for(ri k = j, kp = 0; k < i + j; k ++, kp += kj) {
        int x = a[k], y = mul(w[kp], a[i + k]);
        a[k] = (x + y >= mod) ? x + y - mod : x + y;
        a[i + k] = (x - y < 0) ? x - y + mod : x - y;
    }
}

inline void calc() {
    n = 1; lg = 0;
    while(n <= M + M) n <<= 1, lg ++;
    rep(i, 0, n) rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (lg - 1));
    int g_ = fp(3, (mod - 1) / n);
    w[0] = 1;
    rep(i, 1, n) w[i] = mul(w[i - 1], g_);

    int lim = max(N, n);
    fac[0] = fac[1] = inv[0] = inv[1] = 1;
    rep(i, 2, lim) {
        fac[i] = mul(fac[i - 1], i);
        inv[i] = mul(inv[mod % i], mod - mod / i);
    }
    rep(i, 2, lim) inv[i] = mul(inv[i], inv[i - 1]);

    rep(i, 0, M - 1) f[i] = mul(inv[i], (i & 1) ? mod - 1: 1);
    rep(i, 0, M) if(N >= S * i)
        g[i] = 1ll*fac[i]*C(M,i)%mod*fac[N]%mod*fp(inv[S],i)%mod*inv[N-S*i]%mod*fp(M-i,N-S*i)%mod;
    reverse(g, g + M + 1);

    NTT(f); NTT(g);
    rep(i, 0, n) f[i] = mul(f[i], g[i]);
    NTT(f);
    int ivn = fp(n, mod - 2);
    reverse(f + 1, f + n); reverse(f, f + M + 1);
    rep(i, 0, n) f[i] = mul(f[i], mul(ivn, inv[i]));

    int ans = 0;
    rep(i, 0, M) ans = (ans + mul(f[i], W[i])) % mod;
    printf("%d\n", ans);

}

int main() {
    N = read(); M = read(); S = read();
    rep(i, 0, M) W[i] = read();
    calc();
    return 0;
}

原文地址：https://www.cnblogs.com/reverymoon/p/10173833.html