【CodeForces 997C】Sky Full of Stars（组合计数）

题目链接：【CodeForces 997C】Sky Full of Stars

官方题解：Codeforces Round #493 — Editorial

题目大意：有一个$n\times n$的网格和三种颜色，问有多少种染色方案，使得至少一行或者一列颜色一样。$n\le 10^6$。

如果$n\le 3000$，那么直接暴力容斥即可。
算式：$answer=\sum_{i=0...n,j=0...n,i+j>0}{C_n^iC_n^j(-1)^{i+j+1}3^{(n-i)(n-j)+1}}$

但是因为$n$的数据范围较大，我们考虑将式子变形。
考虑将$i=0\ or\ j=0$的情况拎出来算，时间复杂度$O(n\log n)$。
余下的式子：$answer=\sum_{i=1}^{n}\sum_{j=1}^{n}{C_n^iC_n^j(-1)^{i+j+1}3^{(n-i)(n-j)+1}}$。
换元后得到：$answer=3\sum_{i=0}^{n-1}\sum_{j=0}^{n-1}{C_n^iC_n^j(-1)^{i+j+1}3^{ij}}$。
整理后得到：$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=\sum_{i=0}^n{C_n^ia^ib^{n-i}}$，得到
$answer=3\sum_{i=0}^{n-1}{C_n^i(-1)^{i+1}}[(1+(-3^i))^n-(-3^i)^n]$。

这个式子有$n$项，计算的复杂度为$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;
}