1.题目
问题描述
众所周知，每两周的周三是字节跳动的活动日。作为活动组织者的小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
2.思路
题目含义
首先通过样例一理解题目含义
样例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
nodes = 7：表示有7个节点
decorations = 3：表示有三种礼物类型
tree = [[1, 0, 0, 0, 0, 2, 3], [1, 7], [3, 7], [2, 1], [3, 5], [5, 6], [6, 4]]：
tree[0]— 礼物分配情况
tree[0] = [1, 0, 0, 0, 0, 2, 3] 表示每个节点的礼物类型。具体来说：
tree[0][0] = 1，节点 1 挂有礼物类型 1。tree[0][1] = 0，节点 2 没有挂礼物。tree[0][2] = 0，节点 3 没有挂礼物。tree[0][3] = 0，节点 4 没有挂礼物。tree[0][4] = 0，节点 5 没有挂礼物。tree[0][5] = 2，节点 6 挂有礼物类型 2。tree[0][6] = 3，节点 7 挂有礼物类型 3。
2. tree[1:] — 树的结构
tree[1:] 表示树的边和节点连接关系，其中每个子数组包含一个节点的连接信息。
tree[1] = [1, 7]表示节点 1 和节点 7 相连。tree[2] = [3, 7]表示节点 2 和节点 7 相连。tree[3] = [2, 1]表示节点 3 和节点 1 相连。tree[4] = [3, 5]表示节点 4 和节点 5 相连。tree[5] = [5, 6]表示节点 5 和节点 6 相连。tree[6] = [6, 4]表示节点 6 和节点 4 相连。
树结构示意图
1(1)
/ \\
2(0) 3(0)
/ \\
5(0) 7(3)
/ \\
6(2) 4(2)
要将树划分为 3 个 Special 连通分块，满足以下条件：
- 每个 Special 连通分块内只有一种礼物。
- 每个连通分块可以包含没有礼物的结点。
划分方案
-
划分方案 1：将节点 1 ，2和节点 3 放在一个分块，节点 4 ，5和节点 6 放在另一个分块，节点7放在第三个分块。
-
划分方案 2：节点 1 ，2， 3，5 放在一个分块，节点 7 放在另一个分块，节点 4 和节点 6 放在第三个分块。
- 划分方案 3：节点 1 ，2放在一个分块，节点 3，7 放在另一个分块，节点 4 ，5， 6 放在第三个分块
代码思路
定义is_valid_block函数，判断给定的剪边方案是否能形成有效的连通分块，采用深度优先搜索（DFS），探索所有剪边的可能方案，如果该方案可以形成有效的连通分块，方案数加一
根据边，可以推出树的形状，根据数组可以得到节点上分别挂有礼物类型；要分割为只有单一礼物类型的k个联通量，所以只能砍k-1刀，所以用dfs遍历可能要砍的边，然后剪枝。暴力地用dfs加剪枝做了试试，其中用的Marcode给的邻接表思路表示图，稍微改了一下表示所有相连节点，（这里有个转换为set然后用not in 查询方法的使用）
3.代码
def solution(nodes, decorations, tree):
# Please write your code here
# 定义一个常量 MOD，表示结果需要对其取模 998244353
MOD = 998244353
# 提取礼物信息和边的信息
gifts = tree[0] # gifts 是一个数组，表示每个节点的礼物类型（0表示无礼物，其他数字表示礼物类型）
edges = tree[1:] # edges 是树的边的列表，每一条边是一个由两个节点构成的列表
# 确保 gifts 数组的长度与节点数一致，如果礼物信息不完整，则用 0 填充
if len(gifts) < nodes:
gifts.extend([0] * (nodes - len(gifts)))
# 记录剪边的数组，表示每条边是否被剪掉（n个节点，n-1条边）
cut_edges = [False] * (nodes - 1)
# 定义一个函数，判断给定的剪边方案是否能形成有效的连通分块
def is_valid_block(cut_edges):
# 初始化一个邻接表（邻接表是一个数组，每个元素是一个列表，表示该节点的相邻节点）
# 树节点是从1开始，所以这里空出0号，从1到n
adj = [[] for _ in range(nodes + 1)]
# 遍历所有的边，构建一个邻接表，只有没有被剪掉的边才会被加入到邻接表中
for i in range(nodes - 1): # 边比节点少一条
if cut_edges[i] == False: # 如果这条边没有被剪掉
u, v = edges[i] # 取出边的两个端点 u 和 v
adj_u_set = set(adj[u]) # 对 u 的邻接节点进行去重（用集合存储）
# 将 v 的邻接节点添加到 u 的邻接节点列表中，保持节点之间的连接
# 如果一个节点是v的邻接节点，那么也是u的邻接节点
for w in adj[v]:
if w not in adj_u_set:
adj[u].append(w)
adj_u_set.add(w)
adj[u].append(v)
adj_v_set = set(adj[v]) # 对 v 的邻接节点进行去重
# 将 u 的邻接节点添加到 v 的邻接节点列表中，保持节点之间的连接
for w in adj[u]:
if w not in adj_v_set:
adj[v].append(w)
adj_v_set.add(w)
adj[v].append(u)
# 遍历每个节点，检查该节点所属的块是否符合要求
for i in range(nodes):
# 如果该节点挂有礼物
if gifts[i] != 0:
# 遍历该节点的所有邻接节点，检查是否有礼物类型不同的节点，如果有则返回 False
for node in adj[i + 1]: # node在邻接表中存放的下标为node+1
# gifts[node - 1]表示node节点挂的礼物类型
if gifts[node - 1] != gifts[i] and gifts[node - 1] != 0:
return False
return True
# 定义一个计数器，记录满足条件的划分方案数
count = 0
# 深度优先搜索（DFS）函数，用于探索所有剪边的可能方案
def dfs(i, cut_count):
nonlocal count # 使用 nonlocal 关键字访问外部作用域的 count 变量
# 如果已经剪掉足够的边，或者遍历完所有边，就停止
if cut_count > decorations - 1 or i >= nodes - 1: # cut_count达到decorations - 1即可
return 0 # 终止
# 如果已经剪掉了足够的边，检查当前的剪边方案是否有效
if cut_count == decorations - 1:
if is_valid_block(cut_edges):
count += 1
return 0
# 选择剪掉当前边
cut_edges[i] = True
dfs(i + 1, cut_count + 1) # 递归处理下一个节点，剪掉这条边
# 选择不剪当前边
cut_edges[i] = False
dfs(i + 1, cut_count) # 递归处理下一个节点，不剪这条边
return 0
# 从根节点开始 DFS，初始剪边数为 0
dfs(0, 0)
# 返回符合条件的方案数，对 MOD 取模，防止结果过大
return count % MOD
if __name__ == "__main__":
# You can add more test cases here
testTree1 = [[1,0,0,0,2,3],[1,7],[3,7],[2,1],[3,5],[5,6],[6,4]]
testTree2 = [[1,0,1,0,2],[1,2],[1,5],[2,4],[3,5]]
print(solution(7, 3, testTree1) == 3 )
print(solution(5, 2, testTree2) == 0 )
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
