[CSP-S模拟测试]:长寿花（DP+组合数）

题目描述

庭院里有一棵古树。
圣诞节到了，我想给古树做点装饰，给他一个惊喜。
他会不会喜欢呢？
这棵树可以分为$n$层，第$i$层有$a_i$个防治装饰品的位置，有$m$种颜色的装饰品可供选择。
为了能让他喜欢，我想让装饰品满足以下条件：
$1.$在同一层两个相邻的装饰品不能有同一种颜色；
$2.$相邻两层的颜色集合不能相同（这里颜色集合是指这一层所有出现的颜色去重之后的集合，也就是每个颜色只在集合被最多出现一次）。
我想知道，有多少种不同的方案能让他满意呢？由于方案书可能很多，请告诉我模$p$之后的答案。
谢谢你。

---

输入格式

第一行三个整数，$n,m,p$。
第二行$n$个整数，第$i$个整数表示$a_i$。

---

输出格式

一行一个整数，表示答案。

---

样例

样例输入：

3 2 100
3 1 2

样例输出：

8

---

数据范围与提示

样例解释：

如下集中方法满足条件：
$1.121|1|12$
$2.121|1|21$
$3.121|2|12$
$4.121|2|21$
$5.212|1|12$
$6.212|1|21$
$7.212|2|12$
$8.212|2|21$

数据范围：

记$S=\sum \limits_{i=1}^n a_i$。
对于$20\%$的数据，$1\leqslant m\leqslant 3,S\leqslant 10$。
对于$50\%$的数据，$1\leqslant 1,000,1\leqslant m\leqslant 10$。
对于$100\%$的数据，$1\leqslant n\leqslant {10}^6,1\leqslant m\leqslant {10}^6,1\leqslant a_i\leqslant 5,000,S\leqslant {10}^7,2\leqslant p\leqslant {10}^9$。

---

题解

考虑$DP$。

定义$dp[i][j]$表示到了第$i$层，放了$j$个颜色的装饰品的方案数。

定义$g[i][j]$表示对于一行，到了第$i$个位置，放了$j$个颜色的装饰品的方案数。

我们可以先不管都选了那种颜色，最后在乘上一个组合数就好了。
因为总颜色数是不变的，这样就方便了许多。

$g$数组的转移：

　　$g[i][j]=g[i-1][j-1]*j+g[i-1][j]*(j-1)$。

$dp$数组的转移：

　　$dp[i][j]=dp[i-1][0]*g[a[i]][j]*C[j]-dp[i-1][j]*g[a[i]][j]$

　　$dp[i][0]=dp[i][0]+dp[i][j]$

$dp$数组第一维滚动即可。

时间复杂度：$\Theta(S)$。

期望得分：$100$分。

实际得分：$100$分。

---

代码时刻

#include<bits/stdc++.h>
using namespace std;
int n,m,p;
int a[1000001],maxn,pri[1000001],t[1000001];
bool vis[10000001],pos;
long long dp[2][5001],g[5001][5001],C[5001];
int main()
{
	scanf("%d%d%d",&n,&m,&p);
	for(int i=1;i<=n;i++)
	{
		scanf("%d",&a[i]);
		maxn=max(maxn,a[i]);
	}
	for(int i=2;i<=m;i++)
	{
		if(!vis[i])pri[++pri[0]]=i;
		for(int j=1;j<=pri[0];j++)
		{
			if(i*pri[j]>m)break;
			vis[i*pri[j]]=1;
			if(!(i%pri[j]))break;
		}
	}
	dp[0][0]=g[0][0]=1;
	for(int i=1;i<=maxn;i++)
		for(int j=1;j<=min(m,maxn);j++)
			g[i][j]=(g[i-1][j-1]*j%p+g[i-1][j]*(j-1)%p)%p;
	for(int i=1;i<=min(m,maxn);i++)
	{
		int flag1=i,flag2=m-i+1;
		for(int j=1;j<=pri[0]&&pri[j]<=flag1;j++)
			while(!(flag1%pri[j]))
			{
				flag1/=pri[j];
				t[j]--;
			}
		for(int j=1;j<=pri[0]&&pri[j]<=flag2;j++)
			while(!(flag2%pri[j]))
			{
				flag2/=pri[j];
				t[j]++;
			}
		C[i]=1;
		for(int j=1;j<=pri[0]&&pri[j]<=m;j++)
			for(int k=1;k<=t[j];k++)
				C[i]=(C[i]*pri[j])%p;
	}
	for(int i=1;i<=n;i++)
	{
		pos^=1;
		memset(dp[pos],0,sizeof(dp[pos]));
		for(int j=1;j<=min(m,a[i]);j++)
		{
			dp[pos][j]=(dp[pos^1][0]*g[a[i]][j]%p*C[j]%p-dp[pos^1][j]*g[a[i]][j]%p+p)%p;
			dp[pos][0]=(dp[pos][0]+dp[pos][j])%p;
		}
	}
	printf("%lld",dp[pos][0]);
	return 0;
}

---

rp++