Bytedance Tree 问题｜ 豆包MarsCode AI 刷题

问题背景

在活动日上，字节跳动布置了一棵树形结构的树。树由 n 个结点组成，编号为 1 到 n，树中有 n-1 条边连接这些结点。某些结点挂上了礼物，礼物有 K 种。每个结点只能挂一个礼物，且要求将树划分为若干个 Special 连通分块。每个分块内的结点只能有一种礼物，并且分块可以包含没有礼物的结点。我们的目标是计算出有多少种不同的划分方式。

核心概念

1. 树结构：由 n 个结点和 n-1 条边构成。树是一个无环连通图。
2. 礼物类型：每个结点可能有一个礼物，或没有礼物。没有礼物的结点可以属于任何一个连通块。
3. 连通块：一个连通分块是一个子树，其内的所有结点要么没有礼物，要么拥有相同类型的礼物。

代码解析

数据结构

1. tree: 一个邻接表，用于存储树的结构。
2. gifts: 数组，存储每个结点挂的礼物类型。gifts[i] 表示第 i 个结点的礼物类型，0 表示没有礼物。
3. visited: 布尔数组，用于标记结点是否已被访问。
4. connectedComponents: 记录树中连通块的数量。

核心方法

1. solution 方法

   负责计算和返回树的划分方式数。

   java
   复制代码
   public static int solution(int nodes, int decorations, List<List<Integer>> input) {
       // 初始化树结构和礼物信息
       tree = new ArrayList[nodes + 1];
       gifts = new int[nodes + 1];
       visited = new boolean[nodes + 1];
       connectedComponents = 0;
   
       // 处理礼物信息
       List<Integer> giftInfo = input.get(0);
       for (int i = 1; i <= nodes; i++) {
           gifts[i] = giftInfo.get(i - 1); // 将礼物信息保存
       }
   
       // 构建树的邻接表
       for (int i = 0; i <= nodes; i++) {
           tree[i] = new ArrayList<>();
       }
       for (int i = 1; i < input.size(); i++) {
           List<Integer> edge = input.get(i);
           int u = edge.get(0);
           int v = edge.get(1);
           tree[u].add(v);
           tree[v].add(u);
       }
   
       // 计算树中挂有不同类型礼物的连通块
       for (int i = 1; i <= nodes; i++) {
           if (!visited[i] && gifts[i] != 0) { // 只对挂有礼物的节点进行DFS
               connectedComponents++;
               dfs(i, gifts[i]);
           }
       }
   
       // 如果连通块数量超过礼物种类数，返回 0
       if (connectedComponents > decorations) {
           return 0;
       }
   
       // 结果为连通块的数量，取模 998244353
       return connectedComponents % MOD;
   }
   

   • 初始化: 函数首先初始化了树的邻接表、礼物数组、访问数组和连通块计数器。
   • 构建树的邻接表: 根据输入数据构建树的邻接表。树的边信息保存在输入列表中。
   • DFS 遍历: 通过深度优先搜索 (dfs) 遍历树，找到所有挂有相同类型礼物的连通块。每找到一个连通块，就增加 connectedComponents 的计数。
   • 连通块数量检查: 如果找到的连通块数量大于礼物的种类数 K，返回 0，因为无法将所有连通块分配给礼物种类。
   • 结果计算: 最终返回连通块数量，取模 998244353保证不会超出范围。

2. dfs 方法

   该方法负责遍历树的某一部分，查找挂有相同类型礼物的结点。

   java
   复制代码
   private static void dfs(int node, int giftType) {
       visited[node] = true;
       for (int neighbor : tree[node]) {
           // 如果邻居节点还没有访问过，并且邻居的礼物类型要么为0要么为giftType
           if (!visited[neighbor] && (gifts[neighbor] == 0 || gifts[neighbor] == giftType)) {
               dfs(neighbor, giftType);
           }
       }
   }
   

   • DFS 递归: 从当前节点出发，递归访问所有邻居节点。只有当邻居节点的礼物类型与当前分块的礼物类型相同，或邻居节点没有礼物时，才会继续访问该节点。
   • 标记访问过的节点: 通过 visited 数组，避免重复访问。

总结

通过深度优先搜索 (DFS) 遍历树结构，找到所有连通块并检查是否可以将这些连通块分配给礼物种类。具体的代码逻辑是，先构建树的邻接表，再通过 DFS 查找连通块，最后根据连通块的数量判断是否可以进行合法划分。