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

Условие: дано дерево (т.е. связный неориентированный граф, не имеющий циклов), состоящее из 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]

Решение на языке Go:

Сложность

Временная сложность: O(n)O(n)O(n)

Пространственная сложность: O(n)O(n)O(n)

Код:

func countSubTrees(n int, edges [][]int, labels string) []int {
    ans := make([]int, n)
    adjList := make([][]int, n)

    for _, edge :=  range edges {
        i, j := edge[0], edge[1]
        adjList[i] = append(adjList[i], j)
        adjList[j] = append(adjList[j], i)
    }

    visited := make([]bool, n)
    _ = dfs(0, adjList, ans, labels, visited)
    return ans
}

func dfs(node int, adjList [][]int, ans []int, labels string, visited []bool) []int {
    count := make([]int , 26)
    if visited[node] {
        return count
    }

    visited[node] = true
    for _, child := range adjList[node] {
        cCount := dfs(child, adjList, ans, labels, visited)
        for i := range cCount {
            count[i] += cCount[i]
        }
    }
    
    label := int(labels[node] - 'a')
    count[label]++
    ans[node] = count[label]
    return count
}