洛谷13310 染紫

题目

雪有一棵大小为 n 的树。

雪定义一种树上的染色方案的权值：

设 a 为其红色极大连通块的大小的平方的和。

设 b 为其蓝色极大连通块的大小的平方的和。

这种染色方案的权值为 ab。

树上一些点已经被染上了红或蓝色，请将剩余点分别染成红或蓝色，求所有合法染色方案的权值和。

设待染色节点的个数为 C，则所有合法染色方案共有 2C 个。

答案可能很大，请对 998244353 取模。

做法

连通块size平方求和，很典的套路

假设只有一种颜色求所有连通块Σsize^2，做法是把问题转为 求在连通块里选两个不同点的方案数，这样一个连通块的总计算次数显然是size^2，和原问题结果一样

本题有两种颜色，求的是(Σ红i的size)*(Σ蓝j的size)，拆开之后等于对任意两个红蓝连通块算乘积；一个连通块内部等价于选两个点，那么原问题等价于 对于树上所有染色方案，分别在一个红连通块和一个蓝连通块里各选两个点 的方案数

连通块内部计算总次数为size1^2*size22，所有连通块加起来算得(Σ红i的size)*(Σ蓝j的size)，最后计算时考虑了 选择的块外颜色不同 时 算作不同方案

设\(f[t][0/1][0/1][0/1][0/1]\)，表示子树t内，t颜色选择0/1（color），t同色连通块中1,2号点是否选（d1,d2，当d1=d2时定义变为 已在子树内同色块中选出两个点），以及子树内是否有已在异色块内选出两个点（all）；一个同色块在结算（父亲是异色点）时要求d1=d2，当一个块选出两个点时就会永久保留，合并到异色父亲时变为all_father=1，再异色就变回d1=d2=1

---

设父亲状态为\(f[t][color][d1][d2][all]\)，儿子为\(f[son][color\_son][d1\_son][d2\_son][all\_son]\)，转移为\(f[t][...]*f[son][...] \to f_{new}[t][...]\)

树上dp，分同色异色转移，同色块合并（选重了就非法），异色块要求d1_son=d2_son，考虑父亲儿子中两种颜色的合并（t同色：all_son和d1,d2；t异色：d1_son,d2_son和all）

\(ans=f[1][0][1][1][1]+f[1][1][1][1][1]\)

（设成all主要是因为空间不够不能设r1r2b1b2，同时这样设不用考虑具体颜色，只考虑父亲儿子是否相同

code

dfs爆栈90，于是写成bfs形式

//#pragma comment(linker, "/STACK:10240000000,10240000000")
#include <bits/stdc++.h>
#define fo(a,b,c) for (int a=b; a<=c; a++)
#define fd(a,b,c) for (int a=b; a>=c; a--)
#define add(a,b) a=((a)+(b))%mod
#define mod 998244353
#define ll long long
#define file
using namespace std;

const int N=2e6+10;
int n,T;
int a[N*2][2],ls[N],len;
ll f[N][2][2][2][2];
ll F[2][2][2][2];
int bz[N];
char st[N];

void New(int x,int y)
{
	++len;
	a[len][0]=y;
	a[len][1]=ls[x];
	ls[x]=len;
}

void init(int t)
{
	if (bz[t]!=1)
	{
		fo(r1,0,1)
		fo(r2,0,1)
		f[t][0][r1][r2][0]=1;
	}
	if (bz[t]!=0)
	{
		fo(b1,0,1)
		fo(b2,0,1)
		f[t][1][b1][b2][0]=1;
	}
}

void up(int t,int son)
{
	memset(F,0,sizeof(F));
	fo(c,0,1) fo(d1,0,1) fo(d2,0,1) fo(all,0,1)
	fo(C,0,1) fo(D1,0,1) fo(D2,0,1) fo(All,0,1)
	{
		if (c!=C)
		{
			if (D1==D2)
			{
				if (all==1 && D1==1) continue;
				int newall;
				if (all==0 && D1==1 || all==1 && D1==0) newall=1;
				else newall=0;
				
				if (All==0)
				add(F[c][d1][d2][newall],f[t][c][d1][d2][all]*f[son][C][D1][D2][All]);
				else
				if (d1==0 && d2==0)
				add(F[c][1][1][newall],f[t][c][d1][d2][all]*f[son][C][D1][D2][All]);
			}
		}
		else
		{
			if (d1==1 && D1==1) continue;
			if (d2==1 && D2==1) continue;
			if (all==1 && All==1) continue;
			
			add(F[c][d1|D1][d2|D2][all|All],f[t][c][d1][d2][all]*f[son][C][D1][D2][All]);
		}
	}
	
	memcpy(f[t],F,sizeof(F));
}

//void dfs(int Fa,int t)
//{
//	init(t);
//	for (int i=ls[t]; i; i=a[i][1])
//	if (a[i][0]!=Fa)
//	{
//		dfs(t,a[i][0]);
//		up(t,a[i][0]);
//	}
//}

int d[N],dH,dT;
int faBfs[N];
void bfs()
{
	dH=0,dT=1;
	d[1]=1;
	memset(faBfs,255,sizeof(faBfs));
	faBfs[1]=0;
	while (dH+1<=dT)
	{
		int t=d[++dH];
		for (int i=ls[t]; i; i=a[i][1])
		if (faBfs[a[i][0]]==-1)
		{
			d[++dT]=a[i][0];
			faBfs[a[i][0]]=t;
		}
	}
	
	fd(id,n,1)
	{
		int t=d[id];
		init(t);
		for (int i=ls[t]; i; i=a[i][1])
		if (a[i][0]!=faBfs[t])
		up(t,a[i][0]);
	}
}

void solve()
{
    scanf("%d",&n);
    fo(i,1,n-1)
    {
    	int x,y;
    	scanf("%d%d",&x,&y);
    	New(x,y),New(y,x);
	}
	scanf("%s",st+1);
	fo(i,1,n)
	{
		switch (st[i])
		{
			case 'w':{bz[i]=-1;break;}
			case 'r':{bz[i]=0;break;}
			case 'b':{bz[i]=1;break;}
		}
	}
	
	bfs();
	
	ll ans=0;
	add(ans,f[1][0][1][1][1]);
	add(ans,f[1][1][1][1][1]);
	add(ans,mod);
	printf("%lld\n",ans);
}

int main()
{
//	freopen("T597882.in","r",stdin);

    T=1;
    //scanf("%d",&T);
    for (;T;--T) solve();

    return 0;
}