众所周知，每两周的周三是字节跳动的活动日。作为活动组织者的小A，在这次活动日上布置了一棵 Bytedance Tree。
Bytedance Tree 由 n 个结点构成，每个结点的编号分别为 1，2，3......n，有 n - 1 条边将它们连接起来，根结点为 1。而且为了观赏性，小A 给 M 个结点挂上了 K 种礼物（0 ≤ K ≤ M ≤ N，且保证一个结点只有一个礼物）。
现在小A希望将 Bytedance Tree 划分为 K 个 Special 连通分块，送给参加活动的同学们，请问热心肠的你能帮帮他，告诉小A一共有多少种划分方式吗？
一个 Special 连通分块应该具有以下特性：
- Special 连通分块里只有一种礼物（该种类礼物的数量不限）。
- Special 连通分块可以包含任意数量的未挂上礼物的结点。
由于方案数可能过大，对结果取模 998244353。
测试样例
样例1：
输入：
nodes = 7, decorations = 3, tree = [[1, 0, 0, 0, 0, 2, 3], [1, 7], [3, 7], [2, 1], [3, 5], [5, 6], [6, 4]]
输出：3
样例2：
输入：
nodes = 5, decorations = 2, tree = [[1, 0, 1, 0, 2], [1, 2], [1, 5], [2, 4], [3, 5]]
输出：0
样例3：
输入：
nodes = 6, decorations = 2, tree = [[1, 2, 0, 1, 0, 2], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6]]
输出：0
问题理解
我们需要将一棵树划分为 K 个 Special 连通分块，每个分块中只包含一种礼物。Special 连通分块可以包含任意数量的未挂上礼物的结点。我们需要计算有多少种划分方式，并对结果取模 998244353。
数据结构选择
- 树的表示：使用邻接表来表示树的结构。
- 礼物分布：使用数组来记录每个结点的礼物类型。
- 父结点关系：使用数组来记录每个结点的父结点。
算法步骤
-
初始化：
- 读取礼物分布信息，并记录挂有礼物的结点。
- 检查礼物数量是否与 K 一致，如果不一致则直接返回 0。
- 检查每种礼物类型是否只对应一个结点，如果不一致则直接返回 0。
-
构建树：
- 使用邻接表构建树的结构。
-
确定父结点：
- 使用 BFS 或 DFS 遍历树，确定每个结点的父结点。
-
计算方案数：
- 对于每个挂有礼物的结点，计算从该结点到其最近的挂有礼物的祖先结点之间的路径长度。
- 将这些路径长度的乘积作为总的方案数，并对结果取模 998244353。
关键点
- 路径长度计算：对于每个挂有礼物的结点，计算从该结点到其最近的挂有礼物的祖先结点之间的路径长度。这个路径长度表示了在该路径上可以有多少种划分方式。
- 乘积计算：将所有路径长度的乘积作为总的方案数
关键点
-
未挂上礼物的结点可以被包含在任何 Special 连通分块中：
- 这意味着未挂上礼物的结点可以自由地与挂有相同礼物的结点组合在一起，形成一个 Special 连通分块。
-
路径长度计算：
- 在计算从挂有礼物的结点到其最近的挂有礼物的祖先结点之间的路径长度时，路径上的未挂上礼物的结点也会被计算在内。
具体步骤
-
初始化：
- 读取礼物分布信息，并记录挂有礼物的结点。
- 检查礼物数量是否与 K 一致，如果不一致则直接返回 0。
- 检查每种礼物类型是否只对应一个结点，如果不一致则直接返回 0。
-
构建树：
- 使用邻接表构建树的结构。
-
确定父结点：
- 使用 BFS 或 DFS 遍历树，确定每个结点的父结点。
-
计算方案数：
- 对于每个挂有礼物的结点，计算从该结点到其最近的挂有礼物的祖先结点之间的路径长度。
- 路径长度包括路径上的所有结点，无论它们是否挂有礼物。
- 将这些路径长度的乘积作为总的方案数，并对结果取模 998244353。### 关键点
-
未挂上礼物的结点可以被包含在任何 Special 连通分块中：
- 这意味着未挂上礼物的结点可以自由地与挂有相同礼物的结点组合在一起，形成一个 Special 连通分块。
-
路径长度计算：
- 在计算从挂有礼物的结点到其最近的挂有礼物的祖先结点之间的路径长度时，路径上的未挂上礼物的结点也会被计算在内。
具体步骤
-
初始化：
- 读取礼物分布信息，并记录挂有礼物的结点。
- 检查礼物数量是否与 K 一致，如果不一致则直接返回 0。
- 检查每种礼物类型是否只对应一个结点，如果不一致则直接返回 0。
-
构建树：
- 使用邻接表构建树的结构。
-
确定父结点：
- 使用 BFS 或 DFS 遍历树，确定每个结点的父结点。
-
计算方案数：
- 对于每个挂有礼物的结点，计算从该结点到其最近的挂有礼物的祖先结点之间的路径长度。
- 路径长度包括路径上的所有结点，无论它们是否挂有礼物。
- 将这些路径长度的乘积作为总的方案数，并对结果取模 998244353。
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
public class Main {
public static int solution(int nodes, int decorations, List<List<Integer>> tree) {
// 模拟输入的第一行是礼物分布
List<Integer> decorationsList = tree.get(0);
// 找到所有挂有礼物的结点
List<Integer> decorationNodes = new ArrayList<>();
for(int i = 0; i < decorationsList.size(); i++) {
if(decorationsList.get(i) != 0){
decorationNodes.add(i+1); // 1-based indexing
}
}
// 检查礼物的数量是否与K一致
if(decorationNodes.size() != decorations){
return 0;
}
// 检查每种礼物类型是否只对应一个结点
int[] typeCount = new int[decorations +1];
for(int i = 0; i < decorationsList.size(); i++){
if(decorationsList.get(i) != 0){
typeCount[decorationsList.get(i)]++;
}
}
for(int t = 1; t <= decorations; t++){
if(typeCount[t] !=1){
return 0;
}
}
// 选择第一个礼物结点作为根结点
int root = decorationNodes.get(0);
// 构建树的邻接表
List<List<Integer>> adj = new ArrayList<>(nodes +1);
for(int i = 0; i <= nodes; i++){
adj.add(new ArrayList<>());
}
for(int i =1; i < tree.size(); i++){
List<Integer> edge = tree.get(i);
int u = edge.get(0);
int v = edge.get(1);
adj.get(u).add(v);
adj.get(v).add(u);
}
// BFS 来确定每个结点的父结点
int[] parent = new int[nodes +1];
boolean[] isDecoration = new boolean[nodes +1];
for(int node : decorationNodes){
isDecoration[node] = true;
}
parent[root] = -1;
Queue<Integer> queue = new LinkedList<>();
queue.add(root);
while(!queue.isEmpty()){
int u = queue.poll();
for(int v : adj.get(u)){
if(parent[v] == 0 && v != root){
parent[v] = u;
queue.add(v);
}
}
}
// 计算方案数
long mod = 998244353;
long total =1;
for(int i =1; i < decorationNodes.size(); i++){
int node = decorationNodes.get(i);
int count =0;
int current = node;
while(parent[current] != -1){
int p = parent[current];
count++;
if(isDecoration[p]){
break;
}
current = p;
}
total = (total * count) % mod;
}
return (int) total;
}
public static void main(String[] args) {
List<List<Integer>> testTree1 = new ArrayList<>();
testTree1.add(List.of(1, 0, 0, 0, 0, 2, 3));
testTree1.add(List.of(1, 7));
testTree1.add(List.of(3, 7));
testTree1.add(List.of(2, 1));
testTree1.add(List.of(3, 5));
testTree1.add(List.of(5, 6));
testTree1.add(List.of(6, 4));
List<List<Integer>> testTree2 = new ArrayList<>();
testTree2.add(List.of(1, 0, 1, 0, 2));
testTree2.add(List.of(1, 2));
testTree2.add(List.of(1, 5));
testTree2.add(List.of(2, 4));
testTree2.add(List.of(3, 5));
// 运行并验证结果
System.out.println(solution(7, 3, testTree1) == 3);
System.out.println(solution(5, 2, testTree2) == 0);
}
}
