子树中标签相同的节点数

难度: Medium

给你一棵树（即，一个连通的无环无向图），这棵树由编号从 0 到 n - 1 的 n 个节点组成，且恰好有 n - 1 条 edges 。树的根节点为节点 0 ，树上的每一个节点都有一个标签，也就是字符串 labels 中的一个小写字符（编号为 i 的节点的标签就是 labels[i] ）

边数组 edges 以 edges[i] = [a_i, b_i] 的形式给出，该格式表示节点 a_i 和 b_i 之间存在一条边。

返回一个大小为 n 的数组，其中 ans[i] 表示第 i 个节点的子树中与节点 i 标签相同的节点数。

树 T 中的子树是由 T 中的某个节点及其所有后代节点组成的树。

示例 1：

输入：n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], labels = "abaedcd"
输出：[2,1,1,1,1,1,1]
解释：节点 0 的标签为 'a' ，以 'a' 为根节点的子树中，节点 2 的标签也是 'a' ，因此答案为 2 。注意树中的每个节点都是这棵子树的一部分。
节点 1 的标签为 'b' ，节点 1 的子树包含节点 1、4 和 5，但是节点 4、5 的标签与节点 1 不同，故而答案为 1（即，该节点本身）。

示例 2：

输入：n = 4, edges = [[0,1],[1,2],[0,3]], labels = "bbbb"
输出：[4,2,1,1]
解释：节点 2 的子树中只有节点 2 ，所以答案为 1 。
节点 3 的子树中只有节点 3 ，所以答案为 1 。
节点 1 的子树中包含节点 1 和 2 ，标签都是 'b' ，因此答案为 2 。
节点 0 的子树中包含节点 0、1、2 和 3，标签都是 'b'，因此答案为 4 。

示例 3：

输入：n = 5, edges = [[0,1],[0,2],[1,3],[0,4]], labels = "aabab"
输出：[3,2,1,1,1]

提示：

1 <= n <= 10^5
edges.length == n - 1
edges[i].length == 2
0 <= a_i, b_i < n
a_i != b_i
labels.length == n
labels 仅由小写英文字母组成

Submission

运行时间: 285 ms

内存: 82.4 MB

class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
        # Recursive function to traverse the graph and update counts
        def dfs(node, parent):
            results[node] -= label_counts[labels[node]]
            # Decrement the count of the label for this node
            label_counts[labels[node]] += 1
            # Traverse adjacent nodes (children)
            for adjacent in graph[node]:
                if adjacent != parent:
                    dfs(adjacent, node)
            # Increment the count of the label for this node to include itself
            # label_counts[labels[node]] += 1
            # Update the result for this node 
            # (Total count of the label minus the count before visiting the subtree)
            results[node] += label_counts[labels[node]]

        # Create a graph from the edges (adjacency list)
        graph = defaultdict(list)
        for a, b in edges:
            graph[a].append(b)
            graph[b].append(a)
      
        # Counter to keep track of the label frequencies during DFS
        label_counts = Counter()
        # List to store the result for each node
        results = [0] * n
        # Call DFS starting from the root node (0) with no parent (-1)
        dfs(0, -1)

        # After DFS, results list contains the counts for nodes' labels as required
        return results

Explain

这个题解使用了深度优先搜索（DFS）来遍历树。首先，使用边数组构建一个无向图的邻接表表示。定义一个递归函数 `dfs` 来遍历每个节点。在每次递归调用中，首先减少当前节点标签的计数（这是因为我们在返回到父节点时需要正确的计数），然后遍历所有相邻的节点（即子节点），对每个子节点递归调用 `dfs`。在处理完所有子节点后，增加当前节点标签的计数，并更新结果数组，该数组记录了以当前节点为根的子树中与当前节点标签相同的节点数。整个过程从根节点 0 开始，无父节点调用 `dfs`。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
        def dfs(node, parent):
            # 更新当前节点标签的计数器在进入子树前
            results[node] -= label_counts[labels[node]]
            # 更新当前节点的标签计数
            label_counts[labels[node]] += 1
            # 遍历当前节点的所有相邻节点，即子节点
            for adjacent in graph[node]:
                if adjacent != parent:
                    dfs(adjacent, node)
            # 更新结果数组，记录与当前节点标签相同的子树节点数
            results[node] += label_counts[labels[node]]
        
        # 使用字典存储邻接表
        graph = defaultdict(list)
        for a, b in edges:
            graph[a].append(b)
            graph[b].append(a)
      
        # 初始化标签计数器和结果数组
        label_counts = Counter()
        results = [0] * n
        # 从根节点开始无父节点地调用DFS
        dfs(0, -1)
        
        return results

Explore

在构建邻接表时，对于每一对边数组中的边（a, b），我在图的邻接表中同时添加了从节点a到节点b的边和从节点b到节点a的边。这样做确保了无向图的性质，即无论从哪个节点开始，都可以通过邻接表找到与其直接相连的其他节点。具体来说，当遍历到边（a, b）时，会在邻接表中的a键下添加b，同时在b键下添加a，这样就体现了无向图中边的双向连接性。

确实，使用递归函数进行深度优先搜索在树高度极大时可能导致栈溢出，特别是在树退化为链表的情况下。为了避免这种情况，可以考虑使用非递归的方法，例如使用栈来模拟递归过程。在非递归方法中，我们手动管理一个栈来存储需要访问的节点及其父节点信息，从而控制深度搜索的流程而不依赖函数调用栈。这样可以减少栈溢出的风险，尤其是在处理非常大的数据结构时更加安全和高效。

在 `dfs` 函数中，处理节点标签计数的逻辑如下： 1. 在递归调用进入一个节点之前，首先将该节点的标签在 `label_counts` 字典中的当前计数减去1（如果存在）。这是为了在回溯到父节点时不包括当前节点的计数。 2. 接着，增加当前节点的标签在 `label_counts` 字典中的计数，反映出这个节点的标签现在被包括在内。 3. 之后，对当前节点的每个子节点递归调用 `dfs` 函数。 4. 在处理完所有子节点后，再次更新 `label_counts` 中当前节点标签的计数，此时这个计数应该包括了所有子树中相同标签的节点总数。 5. 最后，更新结果数组 `results`，记录以当前节点为根的子树中与当前节点标签相同的节点数。这样的操作确保了我们在每个节点的处理过程中都能准确地统计出标签的数量。