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^iC_{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;
}