DFS+栈_ 子树中标签相同的节点数_198场周赛第2题

先贴一下用时

• 我觉得思路还蛮清晰的，有问题欢迎评论区指正
  

• 本例是一个多叉树（树是一种特殊的图，不含回路），遍历方式有两种：深搜，广搜

广搜

• 即层次遍历，本题若用BFS来做，需要在节点入队时 保存分支信息，因为题目要求的子树包含的标签信息。

• 如下图所示
  

• 这里暂时没有想到较好的解决办法 -_-#

• 下面看一下 深搜 的思路：

深搜

主要思路为：

• 因为DFS是对应路径的，这里用哈希表暂存路径信息。

• 遇到相同的标签，用栈来 存储 该路径的节点标签信息，并且不会马上更新祖先节点(DFS先遍历到的节点)，而是在遍历下一条路径时（即遍历到叶子节点 或者 遍历完一个节点的所有子树），pop出栈顶元素，更新ans数组中 栈顶元素的值，并累加到栈的下一个元素。

• 用vis数组 来标记已经遍历过的顶点

• 上面单说 有点抽象，下面以一个实际的例子说明
  

代码

class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
        data = collections.defaultdict(list)
        graph = collections.defaultdict(list)
        for a, b in edges:
            graph[a].append(b)
            graph[b].append(a)
        res = [0] * n
        vis = [False] * n
        def dfs(index):
            if not vis[index]:
                vis[index] = True
                data[labels[index]].append([index, 1])
                for i, one in enumerate(graph[index]):
                    if not vis[one]:  # 未访问
                        dfs(one)
                a, b = data[labels[index]].pop()
                res[a] = b
                if data[labels[index]]:  # 非空
                    data[labels[index]][-1][1] += b
                    
        dfs(0)
        return res

• 时间复杂度：$O(n)$，n为顶点数量（边的数量为$n-1$, 连通无环图 - 树 ）
• 空间复杂度：$O(n)$