AI 刷题 42.Bytedance Tree 问题 题解 | 豆包MarsCode AI刷题

42.Bytedance Tree 问题

• 问题描述

众所周知，每两周的周三是字节跳动的活动日。作为活动组织者的小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。

• 思路解析

题目要求我们将一棵树划分为 K 个 Special 连通分块，每个 Special 连通分块中只包含一种礼物（或者不包含礼物）。我们需要计算出所有可能的划分方式，并对结果取模 998244353。

• 数据结构选择

1. 树的表示：使用邻接表来表示树的结构，这样可以方便地进行 DFS 遍历。
2. 礼物信息：使用一个数组 gifts 来存储每个节点的礼物信息。
3. 动态规划表：使用一个二维数组 dp，其中 dp[node][gift] 表示以 node 为根的子树中，包含 gift 礼物的 Special 连通分块的划分方式数。

• 算法步骤

1. 构建树的邻接表：

   • 遍历输入的边列表，构建树的邻接表。

2. 初始化动态规划表：

   • 初始化 dp 数组，大小为 (nodes + 1) x (decorations + 1)，初始值为 0。

3. DFS 遍历树：

   • 从根节点开始进行 DFS 遍历。
   • 对于每个节点，初始化其 dp 值：如果该节点有礼物，则 dp[node][gift] = 1。
   • 遍历该节点的所有子节点，递归调用 DFS。
   • 在回溯时，合并子节点的 dp 值到当前节点的 dp 值。合并的方式需要考虑如何将子树的划分方式合并到当前节点的划分方式中。

4. 计算总的划分方式数：

   • 遍历根节点的 dp 值，累加所有礼物的划分方式数。
   • 对结果取模 998244353。

• 关键点

1. DFS 遍历：DFS 遍历是解决树形结构问题的常用方法，可以方便地处理树的递归结构。
2. 动态规划：通过动态规划表 dp 来记录每个节点的划分方式数，并在回溯时合并子节点的 dp 值。
3. 合并子树的划分方式：在合并子树的划分方式时，需要考虑如何将子树的划分方式合并到当前节点的划分方式中，这通常涉及到组合问题。

• 代码部分

def solution(nodes, decorations, tree):
    # 构建树的邻接表
    from collections import defaultdict
    tree_adj = defaultdict(list)
    for edge in tree:
        tree_adj[edge[0]].append(edge[1])
        tree_adj[edge[1]].append(edge[0])
    
    # 礼物信息数组
    gifts = tree[0]
    
    # 初始化动态规划表
    dp = [[0] * (decorations + 1) for _ in range(nodes + 1)]
    
    # DFS遍历树
    def dfs(node, parent):
        # 初始化当前节点的dp值
        gift = gifts[node - 1]
        if gift != 0:
            dp[node][gift] = 1
        
        # 遍历子节点
        for child in tree_adj[node]:
            if child == parent:
                continue
            dfs(child, node)
            
            # 更新dp值，考虑子树的划分方式
            # 这里需要根据子树的dp值来更新当前节点的dp值
            # 具体实现需要考虑如何合并子树的划分方式
        
    # 从根节点开始DFS
    dfs(1, -1)
    
    # 计算总的划分方式数
    total_ways = 0
    for gift in range(1, decorations + 1):
        total_ways += dp[1][gift]
    
    return total_ways % 998244353

if __name__ == "__main__":
    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 )