Задача: Количество узлов

Условие: дано дерево (т.е. связный неориентированный граф, не имеющий циклов), состоящее из n узлов с числом от 0 до n - 1 и ровно n - 1 ребер. Корнем дерева является узел 0, и каждый узел дерева имеет метку, которая представляет собой символ нижнего регистра, указанный в строковых метках (т.е. узел с номером i имеет метку labels[i]).

Массив ребер задан на ребрах фермы[i] = [ai, bi], что означает наличие ребра между узлами ai и bi в дереве.

Возвращает массив размера n, где и[i] - количество узлов в поддереве узла земли, которые имеют ту же метку, что и узел i.

Поддерево дерева - это дерево, состоящее из узла в T и всех его дочерних узлов.

Пример:

Ввод: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], labels = "abaedcd"

Вывод: [2,1,1,1,1,1,1]

Объяснение:

Ввод: n = 4, edges = [[0,1],[1,2],[0,3]], labels = "bbbb"

Вывод: [4,2,1,1]

Решение:

Code

class Solution:
    def countSubTrees(self, n: int, edges: List[List[int]], labels: str) -> List[int]:
        adj = [[] for _ in range(n)]

        for a, b in edges:
            adj[a].append(b)
            adj[b].append(a)

        count = [0]*len(string.ascii_lowercase)
        answer = [0]*n

        def dfs(node, parent):
            index = ord(labels[node]) - ord('a')
            previous = count[index]

            count[index] += 1

            for child in adj[node]:
                if child != parent:
                    dfs(child,node)
            answer[node] = count[index] - previous
        dfs(0,-1)
        return answer