青训营X豆包MarsCode 技术训练营第七课 ｜ 豆包MarsCode AI 刷题

代码定义了一个名为 solution 的函数，旨在解决一个特定的问题：给定一棵装饰有不同类型礼物的树，以及一个整数 decorations，表示允许剪掉的边的最大数量，计算有多少种不同的方式可以剪掉这些边，使得剩下的每个连通分块（子树）都只包含一种类型的礼物。最终的结果需要对 998244353 取模。

输入参数

• nodes: 整数，表示树中的节点总数。
• decorations: 整数，表示允许剪掉的边的最大数量。
• tree: 列表，其中第一个元素是一个整数列表，表示每个节点上的礼物类型（0 表示无礼物），后续元素是表示树结构的边列表，每条边由两个节点编号组成。

变量初始化

1. MOD: 一个常量，值为 998244353，用于对最终结果取模。
2. gifts: 从 tree 中提取的礼物信息列表。
3. edges: 从 tree 中提取的边信息列表。
4. cut_edges: 一个布尔列表，用于记录哪些边被剪掉，初始化为 False，表示所有边都未被剪掉。

辅助函数

1. is_valid_block(cut_edges) :

   • 目的: 判断在剪掉指定的边后，剩下的每个连通分块是否都只包含一种类型的礼物。

   • 实现:

     • 首先，根据 cut_edges 列表构建一个新的邻接表 adj，表示剪掉指定边后的树结构。
     • 遍历每个节点，检查其相邻节点的礼物类型。如果当前节点有礼物，且存在相邻节点既有礼物且礼物类型不同，则返回 False。
     • 如果所有检查都通过，则返回 True。

深度优先搜索（DFS）

• dfs(i, cut_count) :

  • 目的: 通过深度优先搜索，尝试所有可能的剪边组合，并计算满足条件的组合数量。

  • 参数:

    • i: 当前考虑的边的索引。
    • cut_count: 当前已经剪掉的边的数量。

  • 实现:

    • 终止条件:

      • 如果已经剪掉的边数超过 decorations，或者已经遍历完所有边，则返回。

    • 剪边尝试:

      • 尝试剪掉当前边（cut_edges[i] = True），然后递归调用 dfs 处理下一条边。
      • 尝试不剪当前边（cut_edges[i] = False），然后递归调用 dfs 处理下一条边。

    • 更新结果:

      • 当 cut_count 等于 decorations - 1 时，表示已经剪掉了 decorations - 1 条边，此时检查当前剪边组合是否满足条件（调用 is_valid_block）。
      • 如果满足条件，则 count 加一。

主函数逻辑

• 初始化 gifts 和 edges。
• 如果 gifts 列表的长度小于 nodes，则用 0 填充至 nodes 长度，表示无礼物的节点。
• 初始化 cut_edges 列表。
• 从根节点（索引 0）开始，调用 dfs 函数进行深度优先搜索。
• 返回 count 对 MOD 取模的结果。

示例解析

1. testTree1:

   • 节点总数为 7，允许剪掉的边数为 3。
   • 树的礼物和边信息分别表示每个节点的礼物类型和树的结构。
   • 通过深度优先搜索，尝试所有可能的剪边组合，最终计算出满足条件的组合数量。

2. tree:

   • 节点总数为 6，允许剪掉的边数为 2。
   • 同样，通过深度优先搜索计算满足条件的组合数量。

复杂度分析

• 时间复杂度: 由于需要尝试所有可能的剪边组合，时间复杂度为 O(2^(nodes-1))，其中 nodes-1 是树中的边数。这是一个指数级的时间复杂度，对于较大的 nodes 值，计算会非常耗时。
• 空间复杂度: 主要由 adj 邻接表、cut_edges 列表和递归调用栈组成。空间复杂度为 O(nodes + edges)，即 O(nodes)，因为 edges 的数量最多为 nodes-1。

优化建议

• 剪枝: 在 dfs 过程中，可以添加更多的剪枝条件，提前终止不满足条件的搜索路径，从而减少不必要的计算。
• 记忆化搜索: 对于已经计算过的状态（即特定的 cut_edges 组合），可以将其结果缓存起来，避免重复计算。
• 动态规划: 尝试将问题转化为动态规划问题，利用子问题的解来构建最终解，从而减少计算量。

这段代码是一个典型的深度优先搜索（DFS）应用，通过递归尝试所有可能的剪边组合，并检查每种组合是否满足条件。尽管其时间复杂度较高，但对于较小规模的输入，仍然是一个有效的解决方案。