UOJ#449 喂鸽子

题意：有n个鸽子，你每秒随机喂一只鸽子，每只鸽子吃k次就饱了。求期望多少秒之后全饱了。n <= 50, k <= 1000。

解：有两种做法。一种直接DP的n²k做法在这。我用的是Min-Max容斥 + NTT优化DP。

先Min-Max容斥，由于鸽子是等价的，所以相当于求m只鸽子期望多少秒之后有一只饱了。

讲不清楚...看题解吧。

注意题解中m个盒子的答案为什么放了i + 1个球，这是因为我们把最后一个球拿出来考虑了。所以唯一放满k个球的盒子的贡献是1 / (k - 1)!

#include <bits/stdc++.h>

const int MO = 998244353;
const int N = 50010;

int n, K, fac[N], inv[N], invn[N];
int f[55][N], g[55][N], A[N * 4], B[N * 4], r[N * 4];

inline int C(int n, int m) {
    if(n < 0 || m < 0 || n < m) return 0;
    return 1ll * fac[n] * invn[m] % MO * invn[n - m] % MO;
}

inline int qpow(int a, int b) {
    int ans = 1;
    while(b) {
        if(b & 1) ans = 1ll * ans * a % MO;
        a = 1ll * a * a % MO;
        b = b >> 1;
    }
    return ans;
}

inline int Inv(int x) {
    return qpow(x, MO - 2);
}

inline void prework(int n) {
    static int R = 0;
    if(R == n) return;
    R = n;
    int lm = 1;
    while((1 << lm) < n) lm++;
    for(int i = 0; i < n; i++) {
        r[i] = (r[i >> 1] >> 1) | ((i & 1) << (lm - 1));
    }
    return;
}

inline void NTT(int *a, int n, int f) {
    prework(n);
    for(int i = 0; i < n; i++) {
        if(i < r[i]) {
            std::swap(a[i], a[r[i]]);
        }
    }
    for(int len = 1; len < n; len <<= 1) {
        int Wn = qpow(3, (MO - 1) / (len << 1));
        if(f == -1) Wn = Inv(Wn);
        for(int i = 0; i < n; i += (len << 1)) {
            int w = 1;
            for(int j = 0; j < len; j++) {
                int t = 1ll * a[i + len + j] * w % MO;
                a[i + len + j] = (a[i + j] - t) % MO;
                a[i + j] = (a[i + j] + t) % MO;
                w = 1ll * w * Wn % MO;
            }
        }
    }
    if(f == -1) {
        int inv = Inv(n);
        for(int i = 0; i < n; i++) {
            a[i] = 1ll * a[i] * inv % MO;
        }
    }
    return;
}

int main() {
    
    scanf("%d%d", &n, &K);
    
    fac[0] = inv[0] = invn[0] = 1;
    fac[1] = inv[1] = invn[1] = 1;
    for(int i = 2; i <= n * K; i++) {
        fac[i] = 1ll * fac[i - 1] * i % MO;
        inv[i] = 1ll * inv[MO % i] * (MO - MO / i) % MO;
        invn[i] = 1ll * invn[i - 1] * inv[i] % MO;
    }
    
    int len = 1;
    while(len <= n * K) len <<= 1;
    g[0][0] = 1;
    memcpy(B, invn, K * sizeof(int));
    NTT(B, len, 1);
    for(int i = 1; i <= n; i++) {
        /*for(int j = 0; j <= i * K; j++) {
            /// g[i][j] f[i][j] 
            for(int k = 0; k < K && k <= j; k++) {
                g[i][j] = (g[i][j] + 1ll * g[i - 1][j - k] * invn[k] % MO) % MO;
                f[i][j] = (f[i][j] + 1ll * f[i - 1][j - k] * invn[k] % MO) % MO;
            }
            if(j >= K) f[i][j] = (f[i][j] + 1ll * g[i - 1][j - K] * invn[K - 1] % MO) % MO;
        }*/
        memcpy(A, g[i - 1], n * K * sizeof(int));
        memset(A + n * K, 0, (len - n * K) * sizeof(int));
        NTT(A, len, 1);
        for(int j = 0; j < len; j++) {
            A[j] = 1ll * A[j] * B[j] % MO;
        }
        NTT(A, len, -1);
        for(int j = 0; j <= i * K; j++) {
            g[i][j] = A[j];
        }
        memcpy(A, f[i - 1], n * K * sizeof(int));
        memset(A + n * K, 0, (len - n * K) * sizeof(int));
        NTT(A, len, 1);
        for(int j = 0; j < len; j++) {
            A[j] = 1ll * A[j] * B[j] % MO;
        }
        NTT(A, len, -1);
        for(int j = 0; j <= i * K; j++) {
            f[i][j] = A[j];
            if(j >= K) {
                f[i][j] = (f[i][j] + 1ll * g[i - 1][j - K] * invn[K - 1] % MO) % MO;
            }
        }
    }
    
    int ans = 0;
    for(int i = 1; i <= n; i++) {
        int temp = 0;
        for(int j = K - 1; j <= i * K; j++) {
            temp = (temp + 1ll * fac[j - 1] * f[i][j] % MO * qpow(inv[i], j) % MO * j % MO) % MO;
        }
        temp = 1ll * temp * n % MO * inv[i] % MO * C(n, i) % MO;
        if(i & 1) {
            ans = (ans + temp) % MO;
        }
        else {
            ans = (ans - temp) % MO;
        }
    }
    printf("%d\n", (ans + MO) % MO);
    return 0;
}