CF111D Petya and Coloring

原题传送门

分析

考虑将竖直线从左往右扫，发现左半部分的颜色种类数不减，而右半部分颜色种类数不增。于是如果要保证任意时刻两半颜色数量一样，则这个相等的颜色数量一定不变。考虑竖直线在 $1$、$2$ 列之间和在 $m-1$、$m$ 列之间时，发现这两个时刻左半部分颜色数量相等，也就是 $2$ 到 $m-1$ 列中的所有颜色都在第 $1$ 列中出现过。同理 $2$ 到 $m-1$ 列中的颜色也都在第 $m$ 列出现过。也就是中间可以染的颜色集合一定是两边两列颜色集合之交的子集。其次两边两列的颜色种类数一定相同，因为他们都等于竖直线在 $m-1$、$m$ 列之间时线左边的颜色种类数。

于是可以考虑枚举左右两列颜色的交集大小，然后枚举左右两列放多少种颜色。
先定义 $f[i][j]$ 表示在 $i$ 个位置上按顺序放 $j$ 种颜色的方案数，这个可以 $dp$。设当前交集大小为 $i$，左右两列放 $j$ 种颜色，则中间列的染色方案共有 $i^{n(m-2)}$种。交集里有 $i$ 种颜色，于是有 $C_k^i$ 种选法。两边两列去掉交集后颜色互不相同，于是在剩下的颜色中选出 $2(j-i)$ 种，有 $C_{k-i}^{2(j-i)}$ 种选法。在 $2(j-i)$ 种颜色中选出 $j-i$ 种放在左 / 右边，有 $C_{2(j-i)}^{j-i}$ 种选法。于是 $C_k^i{C}_{k-i}^{2(j-i)}{C}_{2(j-i)}^{j-i}$ 即为 $j$ 种颜色之中恰好有 $i$ 种颜色相同的方案数。两边的 $j$ 种颜色分别有 $f[n][j]\times j!$ 种方案，于是合一起两边共有 $(f[n][j]\times j!{)}^2$ 种方案。最终答案即为 $\sum _{i=0}^n\sum _{j=i}^n(f[n][j]\times j!{)}^2\times C_k^i{C}_{k-i}^{2(j-i)}{C}_{2(j-i)}^{j-i}\times i^{n(m-2)}$。

代码

#include <iostream>
#define int long long
using namespace std;
const int p = 1e9 + 7;
int n, m, k;
int frac[1000005], inv[1000005], invf[1000005];
int C(int a, int b) { return a < b ? 0 : (frac[a] * invf[b] % p * invf[a - b]) % p; }
int f[1005][1005];
int qpow(int x, int y) {
    int ret = 1;
    while (y) {
        if (y & 1) 
            ret = ret * x % p;
        y >>= 1;
        x = x * x % p;
    }
    return ret;
}
signed main() {
    cin >> n >> m >> k;
    if (m == 1) {
        cout << qpow(k, n);
        return 0;
    }
    frac[1] = invf[1] = inv[1] = invf[0] = inv[0] = frac[0] = 1;
    for (int i = 2; i <= max(k, n << 1); i++) inv[i] = (p - p / i) * inv[p % i] % p;
    for (int i = 2; i <= max(k, n << 1); i++) frac[i] = frac[i - 1] * i % p, invf[i] = invf[i - 1] * inv[i] % p;
    int ans = 0;
    f[0][0] = 1;
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= min(i, k); j++) 
            f[i][j] = (f[i - 1][j - 1] + f[i - 1][j] * j % p) % p;
    }
    for (int i = 0; i <= n; i++) {
        int t = qpow(i, n * (m - 2));
        for (int j = i; j <= n; j++) {
            ans = (ans + f[n][j] * frac[j] % p * f[n][j] % p * frac[j] % p * 
                t % p * C(k, i) % p * C(k - i, 2 * (j - i)) % p * C(2 * (j - i), j - i)) % p;
        }
    }
    cout << (ans % p + p) % p;
    return 0;
}