KAJIMA CORPORATION CONTEST 2024（AtCoder Beginner Contest 340） G. Leaf Color（虚树+dp）

题目

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;
}