CF140E New Year Garland

题意

有 \(m\) 种小球，用这些小球装饰一棵 \(n\) 层的圣诞树，每层需要放置 \(a_i\) 个小球。在每一层中，相邻球颜色不同，且相邻两层球颜色集合不同，求装饰圣诞树的方案数，答案对 \(p\) 取模。

\(1\le n,m\le10^6,2\le p\le10^9,1\le a_i\le5000,\sum_{i = 1}^n a_i \le 10^7 \qquad \text{5s,250MB}\)

题解

贺 Daniel_lele 的题解。

先解决每一行的问题，每一行的要求是相邻球颜色不同，于是设 \(f_{i,j}\) 表示 \(i\) 个位置，用 \(j\) 种颜色的小球去填的方案数。状态转移就考虑是使用用过的颜色，还是没用过的颜色。

\[f_{i,j} = f_{i - 1,j}\times(j - 1) + f_{i - 1,j - 1} \]

接下来考虑行与行之间的关系，若两行放的球颜色种数不同，那么集合更不可能相同；若颜色数相同，只需要在选颜色的时候方案减掉一种即可。设 \(g_{i,j}\) 表示第 \(i\) 层放 \(j\) 种颜色的球的方案数。由上述内容，首先枚举当前层颜色种数 \(j\)，上一层颜色种数 \(k\)，若 \(j \neq k\)，则：\(g_{i,j} = f_{a_i,j}\times g_{i - 1,k} \times {m\choose j} \times j!\)，反之 \(g_{i,j} = f_{a_i,j}\times g_{i - 1,k} \times \left({m\choose j - 1} - 1\right)\times j!\)。

\(k\) 的枚举是不必要的，记 \(h_i = \sum_{j = 1}^{a_i} g_{i,j}\)，整理得到：

\[g_{i,j} = f_{a_i,j}\times (h_{i - 1}\times {m\choose j} - g_{i - 1,j})\times j! \]

时间复杂度为 \(\mathcal{O}(n + \sum a)\)。

---

代码实现中，需要注意，因为模数不一定为质数，所以不能使用逆元求组合数，注意边界设为 \(1\)。

#include <bits/stdc++.h>
using namespace std;

const int N = 5e3 + 5,M = 1e6 + 5;

#define int long long

int n,m,P;

int fac[N],C[N];

int a[M],L;

int f[N][N],g[N],_g[N];

void init(){
    fac[0] = 1;
    for (int i = 1;i < N;i++) fac[i] = 1ll * fac[i - 1] * i % P;
    C[0] = 1;
    for (int i = 1;i <= min(m,L);i++) C[i] = 1ll * C[i - 1] * (m - i + 1) % P;
}

signed main(){
    cin >> n >> m >> P;
    for (int i = 1;i <= n;i++)
        cin >> a[i], L = max(L,a[i]);
    init();
    f[0][0] = 1;
    for (int i = 1;i <= L;i++)
        for (int j = 0;j <= min(m,i);j++){
            if (j > 1) (f[i][j] += f[i - 1][j] * (j - 1)) %= P;
            (f[i][j + 1] += f[i - 1][j]) %= P;
        }
    _g[0] = 1;
    for (int i = 1;i <= n;i++){
        int s = 0;
        for (int j = 0;j <= a[i - 1];j++)
            (s += _g[j]) %= P;
        for (int j = 0;j <= a[i];j++){
            g[j] = s * C[j] % P * f[a[i]][j] % P;
            if (j <= a[i - 1]) (g[j] += P - _g[j] * fac[j] % P * f[a[i]][j] % P) %= P;
        }
        for (int j = 0;j <= a[i];j++) _g[j] = g[j];
    }
    int ans = 0;
    for (int i = 0;i <= a[n];i++)
        (ans += _g[i]) %= P;
    cout << ans;
    return 0;
}