Number of Nodes in the Sub-Tree With the Same Label LeetCode solution in Python

1
class Solution:
2
    """
3
    we can solve this using dfs based approach
4
    identify the root of the tree and start a dfs from there
5
    """
6

7
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
8
        graph_map = {i: set() for i in range(n)}
9
        for edge in edges:
10
            graph_map[edge[0]].add(edge[1])
11
            graph_map[edge[1]].add(edge[0])
12

13
        result = [None for _ in range(n)]
14

15
        visited = set()
16

17
        def dfs(index):
18
            visited.add(index)
19
            temp = [0 for _ in range(26)]
20
            temp[ord(labels[index]) - 97] += 1
21
            for idx in graph_map[index]:
22
                if idx not in visited:
23
                    x = dfs(idx)
24
                    temp = [a + b for a, b in zip(temp, x)]
25
            result[index] = temp[ord(labels[index]) - 97]
26
            return temp
27

28
        dfs(0)
29
        return result