Luogu5405 CTS2019氪金手游（容斥原理+树形dp）

Posted 2021-12-02 gloid

　　考虑外向树怎么做。显然设f[i][j]为i子树中出现权值和为j的合法方案的概率，转移做树形背包即可。

　　如果树上只有一条反向边，显然可以先不考虑该边计算概率，再减去将整棵树看做外向树的概率。于是考虑容斥，进一步拓展到多条反向边，就是考虑0条反向边的概率-考虑1条反向边的概率+考虑2条反向边的概率……容斥可以在dp中完成，即遇到反向边时分是否考虑它转移，若考虑乘上-1的系数。

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define N 3010
#define P 998244353
char getc(){char c=getchar();while ((c<‘A‘||c>‘Z‘)&&(c<‘a‘||c>‘z‘)&&(c<‘0‘||c>‘9‘)) c=getchar();return c;}
int gcd(int n,int m){return m==0?n:gcd(m,n%m);}
int read()
{
	int x=0,f=1;char c=getchar();
	while (c<‘0‘||c>‘9‘) {if (c==‘-‘) f=-1;c=getchar();}
	while (c>=‘0‘&&c<=‘9‘) x=(x<<1)+(x<<3)+(c^48),c=getchar();
	return x*f;
}
int n,a[N][4],p[N],size[N],I[N],ans,t;
int d[N][N],f[N][N],h[N];
bool flag[N];
int ksm(int a,int k)
{
	int s=1;
	for (;k;k>>=1,a=1ll*a*a%P) if (k&1) s=1ll*s*a%P;
	return s;
}
int inv(int a){return ksm(a,P-2);}
void inc(int &x,int y){x+=y;if (x>=P) x-=P;}
struct data{int x,y,op;
}e[N];
struct data2{int to,nxt,op;
}edge[N];
void addedge(int x,int y,int op){t++;edge[t].to=y,edge[t].nxt=p[x],edge[t].op=op,p[x]=t;}
void dfs(int k)
{
	flag[k]=1;
	for (int i=1;i<=n;i++)
	if (d[k][i]&&!flag[i]) d[k][i]=d[i][k]=k,dfs(i);
}
void dp(int k)
{
	f[k][0]=1;
	for (int i=p[k];i;i=edge[i].nxt)
	{
		dp(edge[i].to);
		for (int x=size[k];x>=0;x--) h[x]=f[k][x];
		for (int x=0;x<=size[k]+size[edge[i].to];x++) f[k][x]=0;
		if (edge[i].op==0)
		{
			for (int x=size[k];x>=0;x--)
				for (int y=size[edge[i].to];y>=0;y--)
				inc(f[k][x+y],1ll*h[x]*f[edge[i].to][y]%P);
		}
		else
		{
			for (int x=size[k];x>=0;x--)
				for (int y=size[edge[i].to];y>=0;y--)
				inc(f[k][x+y],1ll*(P-1)*h[x]%P*f[edge[i].to][y]%P),
				inc(f[k][x],1ll*h[x]*f[edge[i].to][y]%P);
		}
		size[k]+=size[edge[i].to];
	}
	for (int x=size[k];x>=0;x--) h[x]=f[k][x];
	for (int x=0;x<=size[k]+3;x++) f[k][x]=0;
	for (int x=size[k];x>=0;x--)
		for (int y=3;y>=1;y--)
		inc(f[k][x+y],1ll*h[x]*a[k][y]%P*y*I[x+y]%P);
	size[k]+=3;
}
int main()
{
#ifndef ONLINE_JUDGE
	freopen("a.in","r",stdin);
	freopen("a.out","w",stdout);
	const char LL[]="%I64d\n";
#else
	const char LL[]="%lld\n";
#endif
	n=read();
	for (int i=1;i<=n;i++)
	{
		for (int j=1;j<=3;j++) a[i][j]=read();
		int p=inv(a[i][1]+a[i][2]+a[i][3]);
		for (int j=1;j<=3;j++) a[i][j]=1ll*a[i][j]*p%P;
	}
	for (int i=1;i<n;i++) e[i].x=read(),e[i].y=read(),d[e[i].x][e[i].y]=d[e[i].y][e[i].x]=1;
	dfs(1);
	for (int i=1;i<n;i++)
	{
		if (d[e[i].x][e[i].y]==e[i].y) swap(e[i].x,e[i].y),e[i].op=1;
		addedge(e[i].x,e[i].y,e[i].op);
	}
	for (int i=1;i<=n*3;i++) I[i]=inv(i);
	dp(1);
	for (int i=0;i<=3*n;i++) inc(ans,f[1][i]);
	cout<<ans;
	return 0;
}