AtCoder Beginner Contest 340

比赛链接，官方题解

ID	Difficulty (0-5)	Topic	Status	Code
A	0	Brute Force	Solved	A.cpp
B	0	Brute Force	Solved	B.cpp
C	0	DFS / Math	Solved	C.cpp
D	1	Shortest Path	Solved	D.cpp
E	2	Segment Tree	Solved	E.cpp
F	1	Extended Euclidean algorithm	Solved	F.cpp
G	3	Auxiliary Tree, Dynamic Programming	solved	G.cpp

G - Leaf Color

给定一棵无根树 $T$ ，每个点 $u$ 有一个颜色 $a_u$ 。

问这个树有多少个子图 $T’$ ，满足：

• $T’$ 是一棵树，且 $T’$ 中每个叶节点（度为 $1$）的颜色都相同。

颜色相同想到虚树，枚举 $T’$ 叶子的颜色 $C$ ，$T’$ 一定是这个颜色的虚树的子图。

动态规划计算每个点 $u$ 作为 $T’$ 中最高点时有 $f_u$ 种合法方案（不考虑对 $u$ 的约束）。

并一同计算 $f_u$ 中有 $g_u$ 种方案满足 $u$ 是叶子（$u$ 只选了一个儿子）。
$$
f_u=\prod_{v\in son_u} (f_v+1) - [color_u=C], ~~ g_u=\sum_{v\in son_u} f_v
$$
如果点 $u$ 是当前枚举的颜色（可以做叶子），向答案贡献 $f_u$ 。

否则不能出现 $g_u$ 中的情况， 向答案贡献 $f_u-g_u$ 。

复杂度 $O(n\log n)$ 。