CF 140E. New Year Garland

New Year Garland

题意

用 \(m\) 种颜色去装点 \(n\) 层的圣诞树。圣诞树的第 \(i\) 层恰好由 \(li_i\) 个彩灯串成一行，且同一层内的相邻彩灯颜色不同，同时相邻两层所使用的彩灯的颜色集合不同。问有多少种装点方案，答案对 \(p\) 取模。

分析

计数DP，用 \(f(i, j)\) 表示用 \(j\) 种颜色装点 \(i\) 个彩灯的合法方案数量，对第 \(i\) 个彩灯，可以用新的颜色涂，或者用旧的颜色，显然有：

\[f(i, j) = f(i-1, j-1) + f(i-1, j) \times (j-1) \]

设 \(dp(i, j)\) 表示把前面 \(i\) 层涂好，且第 \(i\) 行用了 \(j\) 种颜色的合法方案数量。

\(sum(i)\) 表示把前面 \(i\) 行涂好的方案数量。

如果不考虑相邻行颜色集重复问题，那么有：

\[dp(i, j) = sum(i-1) * f(li(i), j) * A(m, j) \]

A(m, j) 表示 \(m\) 的 \(j\) 子集的排列。

考虑相邻行颜色集重复问题，那么有：

\[dp(i, j) = dp(i, j) - dp(i-1, j) * f(li(i), j) * j! \]

考虑占用空间比较大，采用滚动数组。

Code

#include <iostream>
#include <algorithm>
using namespace std;

const int N = 5050;

typedef long long ll;

int li[1000010]; // 第i行的彩球数量
int f[N][N]; // f(i, j) 表示用j种颜色涂i个彩球的合法方案数量
int dp[2][N]; // dp(i, j) 表示将前i行涂好，且第i行使用j种颜色涂的方案数量
int sum[2]; // sum(i) = \sum dp[i][j] , 1 <= j <= li 表示所有把前面i行涂好的方案数
int A[N], fac[N];

int main ()
{
    int n, m, p; cin >> n >> m >> p;
    for (int i = 1; i <= n; i ++ ) cin >> li[i];
    
    f[0][0] = 1; // 临界状态
    for (int i = 1; i <= 5000; i ++ )
        for (int j = 1; j <= i; j ++ )
            f[i][j] = (f[i-1][j-1] + (ll)f[i-1][j] * (j-1) % p) % p;
    
    fac[0] = 1; for (int i = 1; i <= 5000; i ++ ) fac[i] = (ll)i * fac[i-1] % p;
    
    A[0] = 1; // 预处理A(m, i)
    for (int i = 1; i <= 5000; i ++ ) A[i] = (ll)A[i-1] * (m - i + 1) % p;
    
    sum[0] = 1; // 给空的彩灯涂，1种方案
    for (int i = 1; i <= n; i ++ ) // 枚举每一行
    {
        sum[i%2] = 0;
        for (int j = 1; j <= li[i]; j ++ ) // 枚举使用的颜色数量
        {
            dp[i%2][j] = (ll)A[j] * sum[1-i%2] % p * f[li[i]][j] % p;
            if (j <= li[i-1])
                dp[i%2][j] = (dp[i%2][j] - (ll)fac[j] * dp[1-i%2][j] % p * f[li[i]][j] % p + p) % p;
            sum[i%2] = (sum[i%2] + dp[i%2][j]) % p;
        }
    }
    cout << sum[n%2] << '\n';
    return 0;
}