ABC340G Leaf Color

题意

给定一棵树 $T$，包含 $n$ 个节点，每个节点有颜色。

求有多少个 $T$ 的导出子图 $T^{\prime }$，满足 $T^{\prime }$ 中的叶子节点颜色相同。答案对 $998244353$ 取模。

$n\le 2\times {10}^5$。

分析

由于叶子节点的限制极其特殊，考虑从叶子的角度思考问题。

如果知道了叶子节点的集合，那么 $T^{\prime }$ 唯一确定。那么我们只需要考虑满足条件的叶子节点的点集数。

考虑对每种颜色求解答案，设 $f_{x,0/1/2}$ 表示 $x$ 子树内选了零个、一个、两个及以上个儿子子树内的点的方案数，初始化 $f_{x,0}=[\mathrm{当}\mathrm{前}\mathrm{点}\mathrm{的}\mathrm{颜}\mathrm{色}\mathrm{是}\mathrm{否}\mathrm{是}\mathrm{枚}\mathrm{举}\mathrm{的}\mathrm{颜}\mathrm{色}]$ 转移考虑是否选择该子树的点以及该点颜色是否是当前考虑的颜色，$f_{x,2}=f_{x+2}+(f_{x+1}+f_{x+2})\cdot g_u,f_{x,1}=f_{x,1}+g_u$，这里 $g_u=f_{u,0}+f_{u,1}+f_{u,2}$。

考虑对每个点计算答案，若选了零个或一个点，那么点 $x$ 也需要被选；若选了两个及以上个点，那么点 $x$ 不能被选。若该点的颜色不等于当前枚举的颜色，则 $f_{x,0},f_{x,1}$ 不能被算进答案；否则，$f_{x,0},f_{x,1},f_{x,2}$ 都可以被算进答案。

时间复杂度 $O(n^2)$，但是发现枚举一个颜色时，转移实际有用的点只有那些“关键点”（即颜色是当前枚举的颜色的点）及其 LCA，所以对每个颜色建虚树即可，$O(n\log n)$。

点击查看代码

//#pragma comment(linker, "/stack:200000000")
//#pragma GCC optimize("Ofast")
//#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,tune=native")
#include<iostream>
#include<cstdio>
#include<cstring>
#include<string>
#include<algorithm>
#include<cmath>
#include<map>
#include<unordered_map>
#include<vector>
#include<queue>
#include<bitset>
#include<set>
#include<ctime>
#include<random>
#define x1 xx1
#define y1 yy1
#define IOS ios::sync_with_stdio(false)
#define ITIE cin.tie(0);
#define OTIE cout.tie(0);
#define FlushIn fread(Fread::ibuf,1,1<<21,stdin)
#define FlushOut fwrite(Fwrite::obuf,1,Fwrite::S-Fwrite::obuf,stdout)
#define PY puts("Yes")
#define PN puts("No")
#define PW puts("-1")
#define P__ puts("")
#define PU puts("--------------------")
#define popc __builtin_popcount
#define pii pair<int,int>
#define mp make_pair
#define fi first
#define se second
#define gc getchar
#define pc putchar
#define pb emplace_back
#define rep(a,b,c) for(int a=(b);a<=(c);++a)
#define per(a,b,c) for(int a=(b);a>=(c);--a)
#define reprange(a,b,c,d) for(int a=(b);a<=(c);a+=d)
#define perrange(a,b,c,d) for(int a=(b);a>=(c);a-=d)
#define graph(i,j,k,l) for(int i=k[j];i;i=l[i].nxt)
#define lowbit(x) (x&-x)
#define lson(x) (x<<1)
#define rson(x) (x<<1|1)
#define mem(x,y) memset(x,y,sizeof x)
//#define double long double
#define int long long
//#define int __int128
using namespace std;
typedef long long i64;
bool greating(int x,int y){return x>y;}
bool greatingll(long long x,long long y){return x>y;}
namespace Fread {
	const int SIZE=1<<21;
	char ibuf[SIZE],*S,*T;
	inline char getc(){if(S==T){T=(S=ibuf)+fread(ibuf,1,SIZE,stdin);if(S==T)return '\n';}return *S++;}
}
namespace Fwrite{
	const int SIZE=1<<21;
	char obuf[SIZE],*S=obuf,*T=obuf+SIZE;
	inline void flush(){fwrite(obuf,1,S-obuf,stdout);S=obuf;}
	inline void putc(char c){*S++=c;if(S==T)flush();}
	struct NTR{~NTR(){flush();}}ztr;
}
/*#ifdef ONLINE_JUDGE
#define getchar Fread::getc
#define putchar Fwrite::putc
#endif*/
inline int rd(){
	int x=0,f=1;char ch=getchar();
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=x*10+ch-48;ch=getchar();}return x*f;
}
inline void write(int x,char ch='\0'){
	if(x<0){x=-x;putchar('-');}
	int y=0;char z[40];
	while(x||!y){z[y++]=x%10+48;x/=10;}
	while(y--)putchar(z[y]);if(ch!='\0')putchar(ch);
}
bool Mbg;
const int maxn=2e5+5,maxm=4e5+5,inf=0x3f3f3f3f,mod=998244353;
const long long llinf=0x3f3f3f3f3f3f3f3f;
int n,c[maxn];
vector<int>G[maxn],v[maxn];
int dfn[maxn],dfncnt;
void pdfs(int x,int y){
	dfn[x]=++dfncnt;
	for(int u:G[x])if(u^y)pdfs(u,x);
}
int lg[maxn];
namespace LCA{
	int dep[maxn],f[maxn][20];
	void dfs(int x,int y){
		dep[x]=dep[y]+1,f[x][0]=y;
		rep(i,1,lg[dep[x]])f[x][i]=f[f[x][i-1]][i-1];
		for(int u:G[x])if(u^y)dfs(u,x);
	}
	int lca(int x,int y){
		if(dep[x]<dep[y])swap(x,y);
		while(dep[x]^dep[y])x=f[x][lg[dep[x]-dep[y]]-1];
		if(x==y)return x;
		per(i,18,0)if(f[x][i]^f[y][i])x=f[x][i],y=f[y][i];
		return f[x][0];
	}
}
int sta[maxn],tp;
bool cmp(int x,int y){return dfn[x]<dfn[y];}
void add_edge(int x,int y){G[x].pb(y),G[y].pb(x);}
int f[maxn][3],g[maxn],col;
int ans;
void dp(int x,int y){
	f[x][1]=f[x][2]=g[x]=0,f[x][0]=(c[x]==col?1:0);
	for(int u:G[x])if(u^y){
		dp(u,x);
		f[x][2]=(f[x][2]+(f[x][1]+f[x][2])%mod*g[u]%mod)%mod;
		f[x][1]=(f[x][1]+g[u])%mod;
	}
	g[x]=(f[x][0]+f[x][1]+f[x][2])%mod;
	if(c[x]==col)ans=(ans+g[x])%mod;
	else ans=(ans+f[x][2])%mod;
}
void solve_the_problem(){
	n=rd();
	rep(i,1,n){lg[i]=lg[i-1];if((1<<lg[i])==i)lg[i]++;}
	rep(i,1,n)c[i]=rd(),v[c[i]].pb(i);
	rep(i,2,n){int x=rd(),y=rd();add_edge(x,y);}
	LCA::dfs(1,0);
	pdfs(1,0);
	for(col=1;col<=n;++col)if(v[col].size()){
		sort(v[col].begin(),v[col].end(),cmp);
		tp=sta[1]=1,G[1].clear();
		for(int i:v[col])if(i^1){
			int z=LCA::lca(sta[tp],i);
			if(z^sta[tp]){
				while(dfn[z]<dfn[sta[tp-1]])add_edge(sta[tp],sta[tp-1]),--tp;
				if(z==sta[tp-1])add_edge(sta[tp],sta[tp-1]),tp--;
				else G[z].clear(),add_edge(z,sta[tp]),sta[tp]=z;
			}
			G[i].clear(),sta[++tp]=i;
		}
		rep(i,2,tp)add_edge(sta[i-1],sta[i]);
		dp(1,0);
//		write(ans,32);
	}
	write(ans);
}
bool Med;
signed main(){
//	freopen(".in","r",stdin);freopen(".out","w",stdout);
//	fprintf(stderr,"%.3lfMB\n",(&Mbg-&Med)/1048576.0);
	int _=1;while(_--)solve_the_problem();
}
/*
4
1 1 1 1
1 2
1 3
1 4
*/