[CSP-S模拟测试]:chemistry（期望DP+组合数学）

题目传送门（内部题27）

---

输入格式

第一行有$4$个整数$n,k,p,q$。
第二行有$n$个整数$a_i$。
接下来有$n-1$行，每行有两个整数$u,v$，表示$u$与$v$之间通过化学单键连接。

---

输出格式

一行一个整数表示答案。

---

样例

样例输入：

3 2 1 2
1 2 4
1 2
1 3

样例输出：

500000019

---

数据范围与提示

样例解释：

当没有原子被活化时，概率为$\frac{1}{8}$，爆发值为${(1+2+4)}^2=49$。
当$1$号原子被活化时，概率为$\frac{1}{8}$，爆发值为$2^2+4^2=20$。
当$2$号原子被活化时，概率为$\frac{1}{8}$，爆发值为${(1+4)}^2=25$。
当$3$号原子被活化时，概率为$\frac{1}{8}$，爆发值为${(1+2)}^2=9$。
当$1,2$号原子被活化时，概率为$\frac{1}{8}$，爆发值为$4^2=16$。
当$1,3$号原子被活化时，概率为$\frac{1}{8}$，爆发值为$2^2=4$。
当$2,3$号原子被活化时，概率为$\frac{1}{8}$，爆发值为$1^2=1$。
当$1,2,3$号原子被活化时，概率为$\frac{1}{8}$，爆发值为$0$。
期望为$\frac{124}{8}=\frac{31}{2}$，即$500000019$。

提示：

$\frac{a}{b}$在模质数$m$意义下的值为$a\times b^{m-2}$。
例如，$\frac{31}{2}$在模$1000000007$意义下的值为$31\times 2^{1000000005}\equiv 500000019(\mod {10}^9+7)$。

数据范围：

对于所有数据，$n\leqslant 2\times {10}^5,2\leqslant k\leqslant 10,0\leqslant p<q\leqslant {10}^3,0\leqslant a_i\leqslant {10}^3$。

---

题解

为了简化问题，我们可以将点权忽略，设所有点的点权都为$1$。
问题可以转化为，给定一颗$n$个节点的无根树，每个节点被选的概率都为$p$，令$q=1−p$，求被选中的点形成的所有联通块的大小的$K$次方的和的期望。

再来普及一下有关这道题的知识点：

对于任意两个事件$X,Y$：
$E(X+Y)=E(X)+E(Y)$。
对于两个相互独立事件$X,Y$：
$E(XY)=E(X)E(Y)$。
因此，对于两个相互独立事件$X,Y$，有：
$E((X+Y)j)=\sum \limits_{k=0}^j C_j^k E(X^k)E(Y^{j−k})$。

那么现在我们开始考虑如何解这道题了。

设$f[i][j]$表示以$i$为根的子树中的所有联通块的大小的$j$次方的和的期望。

设$g[i][j]$表示以$i$为根的子树中$i$所在的连通块的大小的$j$次方的期望。

考虑初始状态，$f[x][j]=g[x][j]=p$。

先来考虑$g$数组的转移：

　　将$u$节点的状态与它的一个儿子节点$v$的状态合并，令合并后的状态为$g'$。
　　为了方便转移，我们可以先将$g[u][k]$除以$p$表示在$u$节点一定选的情况下的$u$所在连通块的大小的$k$次方的期望，那么$\frac{g[u][k]}{p}$和$g[v][j−k]$可以看作两个相互独立事件的期望，可以根据相互独立事件的期望公式$E({(X+Y)}^j)=\sum \limits_{k=0}^j C_j^k E(X^k)E(Y^{j-k})$进行状态合并。只不过合并后的$g[u][j]'$仍表示在$u$节点一定选的情况下的$u$所在连通块的大小的$j$次方的期望，所以需要再乘以$p$。
$g[u][j]'=p\times \sum \limits_{k=0}^j C_j^k\times \frac{g[u][k]}{p}\times g[v][j−k]$。
即$g[u][j]'=p\times g[v][j]+g[u][j]+\sum \limits_{k=1}^{j-1} C_j^k\times g[u][k]\times g[v][j−k]$。

再来考虑$f$数组的转移：

将$u$节点的状态与它的一个儿子节点$v$的状态合并时，由于$u$的状态和$v$的状态可看做相互独立事件，我们只需要考虑$u$和$v$衔接部分所产生的贡献。
$f[u][j]'=f[u][j]+f[v][j]+\sum \limits_{k=1}^{j−1} C_j^k\times g[u][k]\times g[v][j−k]$。

进行一遍树形$DP$后$f[1][K]$即为最后答案。
对于题目中有点权的情况，我们只需将初始状态设为$f[x][j]=g[x][j]=p\times a_x^j$即可。

---

代码时刻

#include<bits/stdc++.h>
using namespace std;
struct rec
{
	int nxt;
	int to;
}e[400000];
int head[200001],cnt;
int n,K;
long long p,q;
long long a[200001][11];
long long C[11][11];
long long dp[200001][11],g[2][200001][11];
void pre_work()
{
	C[0][0]=1;
	for(int i=1;i<=10;i++)
	{
		C[i][0]=1;
		for(int j=1;j<=i;j++)
			C[i][j]=C[i-1][j]+C[i-1][j-1];
	}
}
void add(int x,int y)
{
	e[++cnt].nxt=head[x];
	e[cnt].to=y;
	head[x]=cnt;
}
long long qpow(long long x,long long y)
{
	long long res=1;
	while(y)
	{
		if(y&1)res=res*x%1000000007;
		x=x*x%1000000007;
		y>>=1;
	}
	return res;
}
void dfs(int x,int fa)
{
	int son=0;
	for(int i=head[x];i;i=e[i].nxt)
		if(e[i].to!=fa)
			dfs(e[i].to,x);
	for(int i=1;i<=K;i++)dp[x][i]=g[0][x][i]=p*a[x][i]%1000000007;
	for(int i=head[x];i;i=e[i].nxt)
	{
		if(e[i].to==fa)continue;
		for(int j=1;j<=K;j++)
		{
			dp[x][j]=(dp[x][j]+dp[e[i].to][j])%1000000007;
			for(int k=1;k<j;k++)
				dp[x][j]=(dp[x][j]+C[j][k]*g[0][x][k]%1000000007*g[0][e[i].to][j-k]%1000000007)%1000000007;
			dp[x][j]=dp[x][j];
		}
		for(int j=1;j<=K;j++)
		{
			g[1][x][j]=(g[0][x][j]+p*g[0][e[i].to][j]%1000000007)%1000000007;
			for(int k=1;k<j;k++)
				g[1][x][j]=(g[1][x][j]+C[j][k]*g[0][x][k]%1000000007*g[0][e[i].to][j-k]%1000000007)%1000000007;
		}
		for(int j=1;j<=K;j++)g[0][x][j]=g[1][x][j];
	}
}
int main()
{
	pre_work();
	scanf("%d%d%lld%lld",&n,&K,&p,&q);
	p=p*qpow(q,1000000005)%1000000007;
	for(int i=1;i<=n;i++)
	{
		scanf("%lld",&a[i][1]);
		for(int j=2;j<=K;j++)
			a[i][j]=(a[i][j-1]*a[i][1])%1000000007;
	}
	for(int i=1;i<n;i++)
	{
		int u,v;
		scanf("%d%d",&u,&v);
		add(u,v);
		add(v,u);
	}
	dfs(1,0);
	printf("%lld",dp[1][K]);
	return 0;
}

---

rp++