1519. Number of Nodes in the Sub-Tree With the Same Label

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

Python

題目怪怪的，說是會給0為root的tree，但test case卻不一定。下面這個test case會被認為2是1的parent node...

4
[[0,2],[0,3],[1,2]]
"aeed"

DFS with hashmap

from collections import defaultdict, Counter


class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
        mapper = defaultdict(set)

        for edge in edges:
            mapper[edge[0]].add(edge[1])
            mapper[edge[1]].add(edge[0])

        def dfs(node, pre, result):
            counter = Counter()
            counter[labels[node]] += 1

            for nei in mapper[node]:
                if nei == pre:
                    continue
                counter += dfs(nei, node, result)
            result[node] = counter[labels[node]]
            return counter

        result = [0] * n

        dfs(0, None, result)

        return result