【CodeForces 997C】Sky Full of Stars（組合計數）

題目鏈接：【CodeForces 997C】Sky Full of Stars

官方題解：Codeforces Round #493 — Editorial

題目大意：有一個n×n$n\times n$ 的網格和三種顏色，問有多少種染色方案，使得至少一行或者一列顏色一樣。n≤106$n\le {10}^6$ 。

如果n≤3000$n\le 3000$ ，那麼直接暴力容斥即可。
算式：answer=∑i=0...n,j=0...n,i+j>0CinCjn(−1)i+j+13(n−i)(n−j)+1$answer=\sum _{i=\mathrm{0...}n,j=\mathrm{0...}n,i+j>0}{C}_n^i{C}_n^j(-1{)}^{i+j+1}{3}^{(n-i)(n-j)+1}$

但是因爲n$n$ 的數據範圍較大，我們考慮將式子變形。
考慮將i=0 or j=0$i=0orj=0$ 的情況拎出來算，時間複雜度O(nlogn)$O(n\log n)$ 。
餘下的式子：answer=∑ni=1∑nj=1CinCjn(−1)i+j+13(n−i)(n−j)+1$answer=\sum _{i=1}^n\sum _{j=1}^n{C}_n^i{C}_n^j(-1{)}^{i+j+1}{3}^{(n-i)(n-j)+1}$ 。
換元后得到：answer=3∑n−1i=0∑n−1j=0CinCjn(−1)i+j+13ij$answer=3\sum _{i=0}^{n-1}\sum _{j=0}^{n-1}{C}_n^i{C}_n^j(-1{)}^{i+j+1}{3}^{ij}$ 。
整理後得到：answer=3∑n−1i=0Cin(−1)i+1∑n−1j=0Cjn(−3i)j$answer=3\sum _{i=0}^{n-1}{C}_n^i(-1{)}^{i+1}\sum _{j=0}^{n-1}{C}_n^j(-3^i{)}^j$ 。

由(a+b)n=∑ni=0Cinaibn−i$(a+b{)}^n=\sum _{i=0}^n{C}_n^i{a}^i{b}^{n-i}$ ，得到
answer=3∑n−1i=0Cin(−1)i+1[(1+(−3i))n−(−3i)n]$answer=3\sum _{i=0}^{n-1}{C}_n^i(-1{)}^{i+1}[(1+(-3^i){)}^n-(-3^i{)}^n]$ 。

這個式子有n$n$ 項，計算的複雜度爲O(nlogn)$O(n\log n)$ 。

#include <cstdio>
const int maxn = 1000005;
const int mod = 998244353;
int res, ans, n;
int fac[maxn], inv[maxn];
void update(int &x, int y) {
    x += y;
    if (x < 0) {
        x += mod;
    }
    if (x >= mod) {
        x -= mod;
    }
}
int mpow(int x, int y) {
    update(x, 0);
    int z = 1;
    for (; y; y >>= 1) {
        if (y & 1) {
            z = 1ll * z * x % mod;
        }
        x = 1ll * x * x % mod;
    }
    return z;
}
int calc(int x, int y) {
    return 1ll * fac[x] * inv[y] % mod * inv[x - y] % mod;
}
int main() {
    scanf("%d", &n);
    fac[0] = 1, inv[0] = 1;
    for (int i = 1; i <= n; i++) {
        fac[i] = 1ll * fac[i - 1] * i % mod;
        inv[i] = mpow(fac[i], mod - 2);
    }
    for (int i = 1; i <= n; i++) {
        int a = 1ll * calc(n, i) * mpow(-1, i + 1) % mod;
        int x = mpow(3, (1ll * n * (n - i) + i) % (mod - 1));
        update(ans, 1ll * a * x % mod);
    }
    ans = 2 * ans % mod;
    for (int i = 0; i < n; i++) {
        int t = mod - mpow(3, i);
        int x = (mpow(t + 1, n) + mod - mpow(t, n)) % mod;
        int a = 1ll * calc(n, i) * mpow(-1, i + 1) % mod;
        update(res, 1ll * a * x % mod);
    }
    res = 3ll * res % mod;
    printf("%d\n", (ans + res) % mod);
    return 0;
}