CodeForces 140 E.New Year Garland（组合数学+dp）

Description

n 排，第 i 排 li 个位置，现在要给所有位置用 m 种颜色上色，需要满足：

1.每一排相邻位置不同色

2.相邻排所用颜色集不同

问方案数，结果模 p

Input

第一行三个整数 n,m,p ，之后输入 n 个整数 li (1≤n,m≤106,2≤p≤109,1≤li≤5000,∑i=1nli≤107)

Output

方案数，结果模 p

Sample Input

3 2 1000
3 1 2

Sample Output

8

Solution

dp[i][j] 表示已经给前 i 排染好色且第 i 排用 j 种颜色的方案数， f[i][j] 表示用 j 种颜色给 i 个位置染色且保证相邻两个位置颜色不同的方案数， sum[i]=∑j=1lidp[i][j] 表示给前 i 排染好色的方案数

简单转移得到 f[i][j] ： f[0][0]=1,f[i][j]=f[i−1][j−1]+(j−1)⋅f[i−1][j]

不考虑相邻两排颜色集不同时， dp[i][j]=sum[i−1]⋅Ajm⋅f[li][j]

如果 i>1 且 j≤li−1 ，那么就存在当前颜色集和上一行颜色集重复的不合法情况， dp[i][j]−=dp[i−1][j]⋅j!⋅f[li][j] ，由于在求 dp[i−1][j] 的时候已经考虑选出的颜色及其编号，故显然只需考虑用这选出的 j 种颜色给 li 个位置染色，故乘了 j!⋅f[li][j]

答案即为 sum[n] ，由于空间限制第一维可以滚动，时间复杂度 O(L)

Code

#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
using namespace std;
typedef long long ll;
typedef pair<int,int>P;
const int INF=0x3f3f3f3f,maxn=1000005;
int n,m,p,l[maxn],A[5005],fact[5005],f[5005][5005],dp[2][5005],sum[2];
int main()
{
    scanf("%d%d%d",&n,&m,&p);
    for(int i=1;i<=n;i++)scanf("%d",&l[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;
    A[0]=1;
    for(int i=1;i<=5000;i++)A[i]=(ll)A[i-1]*(m+1-i)%p;
    fact[0]=1;
    for(int i=1;i<=5000;i++)fact[i]=(ll)i*fact[i-1]%p;
    sum[0]=1;
    for(int i=1;i<=n;i++)
    {
        sum[i%2]=0;
        for(int j=1;j<=l[i];j++)
        {
            dp[i%2][j]=(ll)A[j]*sum[1-i%2]%p*f[l[i]][j]%p;
            if(i&&j<=l[i-1])dp[i%2][j]=(dp[i%2][j]-(ll)fact[j]*dp[1-i%2][j]%p*f[l[i]][j]%p+p)%p;
            sum[i%2]=(sum[i%2]+dp[i%2][j])%p; 
        }
    }
    printf("%d\n",sum[n%2]);
    return 0;
}