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&