Leetcode 1643 Solution

class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
        tree = collections.defaultdict(list)
        for src, dst in edges:
            tree[src].append(dst)
            tree[dst].append(src)

        cache = {}

        def dfs(node_id):
            nonlocal tree
            nonlocal labels
            nonlocal cache

            if node_id in cache:
                return cache[node_id]

            # Fill up the cache first.
            cache[node_id] = None

            # Prepare an agg dict.
            agg_dict = collections.defaultdict(int)

            # Walk through the child nodes, and update the agg dict.
            for child_node_id in tree[node_id]:
                child_agg_dict = dfs(child_node_id)

                if child_agg_dict is None:
                    continue

                for k, v in child_agg_dict.items():
                    agg_dict[k] += v

            # Update the current node.
            agg_dict[labels[node_id]] += 1

            # Generate the result, and return it.
            cache[node_id] = agg_dict

            return agg_dict

        return [dfs(i)[label] for i, label in enumerate(labels)]