题目
n(n<=2e5)个点的树,每个点有一个颜色ai(1<=ai<=n)
求满足以下两个条件的导出子图的方案数,答案对998244353取模:
1. 是一棵树
2. 每个度为1的结点的颜色都相同
导出子图,即在n个点里选一个子集,选出一些点来,保留这些点之间原来的边,得到的图
思路来源
乱搞ac
题解
枚举每种颜色,对每种颜色求虚树,
因为,对于非叶子结点,除了lca是必经的点以外,
其他点要么顺带经过,要么不选,所以无需关注
求完虚树之后,问题转化成,对于某个给定颜色,
对于一棵树,所有叶子结点,都是关键点,
非叶子结点,有一些节点是关键点,有一些不是关键点
在虚树上取一个连通块,使得连通块的叶子都是关键点,求方案数
dp[i][2]表示以i为根,i不取/i取的方案数
1. dp[i][0]表示i不取,i不取的话所有子树都不能取
2. dp[i][1]表示i取,i取的话,子树可取可不取
特别地,当i是非关键点的时候,i不能当叶子,这一种情况需要减掉
枚举以i为根时的连通块的方案数,统计答案
1. 如果i是关键点,dp[i][1]是i为根的贡献
2. 如果i不是关键点,则i需要至少选两棵子树
因为非关键点不能当叶子,所以子树里一定有关键点
f[0/1/2]表示i的子树里选了0/1/>=2棵的方案数,背包考虑子树即可,
此时f[2]是i为根的贡献
代码
#include<bits/stdc++.h>
#include<iostream>
#include<cstdio>
#include<vector>
#include<map>
#include<queue>
using namespace std;
typedef array<int,2> a2;
typedef array<int,3> a3;
const int N=2e5+10,LG=20,mod=998244353;
int n,u,v,ans,a[N];
vector<int>col[N],e[N],g[N];
int dep[N],f[N][LG+1],dfn[N],c;
bool key[N];//关键点
int lca(int x,int y){
if(dep[x]<dep[y])swap(x,y);
for(int i=LG;i>=0;i--){
if(dep[f[x][i]]>=dep[y]){
x=f[x][i];
}
}
if(x==y)return x;
for(int i=LG;i>=0;i--){
if(f[x][i]!=f[y][i]){
x=f[x][i],y=f[y][i];
}
}
//printf("x:%d y:%d lca:%d\n",x,y,f[x][0]);
return f[x][0];
}
void dfs(int u,int fa){
dep[u]=dep[fa]+1;
dfn[u]=++c;
//printf("u:%d dfn:%d\n",u,dfn[u]);
f[u][0]=fa;
for(int i=1;i<=LG;i++){
f[u][i]=f[f[u][i-1]][i-1];
}
for(auto &v:e[u]){
if(v==fa)continue;
dfs(v,u);
}
}
void add(int x,int y){
//if(x==y)return;
//printf("x:%d y:%d\n",x,y);
g[x].push_back(y);
}
a2 dfs1(int u){
//printf("u:%d\n",u);
a2 dp={1,1};
a3 h={1,0,0};
for(auto &v:g[u]){
auto dp2=dfs1(v);
dp[0]=1ll*dp[0]*dp2[0]%mod;
dp[1]=1ll*dp[1]*(dp2[1]+dp2[0])%mod;
h[2]=(1ll*h[2]*(dp2[0]+dp2[1])%mod+1ll*h[1]*dp2[1]%mod)%mod;
h[1]=(1ll*h[1]*dp2[0]%mod+1ll*h[0]*dp2[1]%mod)%mod;
h[0]=1ll*h[0]*dp2[0]%mod;
}
if(key[u]){
ans=(ans+dp[1])%mod;
}
else{
ans=(ans+h[2])%mod;
dp[1]=(dp[1]+mod-1)%mod;//非关键点的不能当叶子
}
key[u]=0;
g[u].clear();
return dp;
}
void build(vector<int>&a){//以0为虚根
static int stk[N],top;//虚树栈
sort(a.begin(),a.end(),[&](int x,int y){
return dfn[x]<dfn[y];
});
int sz=a.size();
stk[top=1]=a[0];
key[a[0]]=1;// 标记为关键点
for(int i=1;i<sz;++i){
int u=a[i];
key[u]=1;// 标记为关键点
int p=lca(u,stk[top]);
if(p!=stk[top]){
while(dep[stk[top-1]]>dep[p]){
add(stk[top-1],stk[top]);
top--;
}
add(p,stk[top]);
top--;
if(stk[top]!=p){
stk[++top]=p;
}
}
stk[++top]=u;
}
for(int i=top;i>=2;i--){
add(stk[i-1],stk[i]);
}
dfs1(stk[1]);
}
int main(){
scanf("%d",&n);
for(int i=1;i<=n;++i){
scanf("%d",&a[i]);
col[a[i]].push_back(i);
}
for(int i=1;i<n;++i){
scanf("%d%d",&u,&v);
e[u].push_back(v);
e[v].push_back(u);
}
dfs(1,0);
for(int i=1;i<=n;++i){
if(!col[i].size())continue;
build(col[i]);
//puts("");
}
printf("%d\n",ans);
return 0;
}