leetcode 1519. 子树中标签相同的节点数

原题链接

https://leetcode-cn.com/problems/number-of-nodes-in-the-sub-tree-with-the-same-label/

原题简要描述

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

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

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

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

这道题目问题的描述有点绕，我第一遍没有读明白，后来才明白。

其实就是将树变成无环无向图，有一个顶点。所有的边都给你了，顶点用数字表示，数字对应一个字符串的索引，让你写出每个数字的子树上（包括这个数字的顶点本身），有多少个和这个数字对应索引的字符相同的。

代码

# 定义一个子函数
# 这个函数计算每一个节点为根的子树，所有标记的计数
# 因为用邻接表表示的边，我们为了防止边的重复访问，需要记录访问节点
# 因为没有环，所以我们只要记住前一个节点，在访问下一个节点的相邻节点的时候
# 不访问这个节点就可以了
# 参数为dfs(i, pre)， i就是这个节点的编号，pre就是前一个节点的编号
# 终止条件就是这个点的相邻节点都访问过了
# 这样我们在全局的一个数组变量中，更新相应的i的值即可
# 最后我们返回这个数组变量就是所求
# 利用了python中Counter有加法的api
# 我们可以简化代码
# 否则我们如果用字典，我们要遍历字典，把存在于本层的加上，不存在于本层的添加进来即可
# 稍微麻烦点
class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> \
            List[int]:
        edge_map = defaultdict(list)
        for e in edges:
            edge_map[e[0]].append(e[1])
            edge_map[e[1]].append(e[0])
        ans = [1] * n

        def dfs(i, pre):
            data = {labels[i]: 1}
            for nxt in edge_map[i]:
                if nxt != pre:
                    for k, v in dfs(nxt, i).items():
                        if k in data:
                            data[k] += v
                        else:
                            data[k] = v
            ans[i] = data[labels[i]]
            return data

        dfs(0, None)
        return ans

class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> \
            List[int]:
        edge_map = defaultdict(list)
        for e in edges:
            edge_map[e[0]].append(e[1])
            edge_map[e[1]].append(e[0])
        ans = [1] * n

        def dfs(i, pre):
            data = Counter({labels[i]: 1})
            for nxt in edge_map[i]:
                if nxt != pre:
                    data += dfs(nxt, i)
            ans[i] = data[labels[i]]
            return data

        dfs(0, None)
        return ans