Сложность: medium
Вам дано дерево (т.е. связный неориентированный граф без циклов), состоящее из n узлов, пронумерованных от 0 до n - 1, и ровно n - 1 ребра. Корнем дерева является узел 0, и каждый узел дерева имеет метку, которая является строчной буквой, указанной в строке labels (т.е. узел с номером i имеет метку labels[i]).
Массив edges дан в форме edges[i] = [ai, bi], что означает, что существует ребро между узлами ai и bi в дереве.
Верните массив размера n, где ans[i] — это количество узлов в поддереве узла i, которые имеют ту же метку, что и узел i.
Поддерево дерева T — это дерево, состоящее из узла в T и всех его дочерних узлов.
Пример:
Input: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], labels = "abaedcd"
Output: [2,1,1,1,1,1,1]
Explanation: Node 0 has label 'a' and its sub-tree has node 2 with label 'a' as well, thus the answer is 2. Notice that any node is part of its sub-tree.
Node 1 has a label 'b'. The sub-tree of node 1 contains nodes 1,4 and 5, as nodes 4 and 5 have different labels than node 1, the answer is just 1 (the node itself).
👨💻 Алгоритм:
1⃣Создайте список смежности, где adj[X] содержит всех соседей узла X.
2⃣Инициализируйте массив ans, хранящий ответ для каждого узла, и заполните его нулями.
3⃣Начните обход в глубину (DFS).
😎 Решение:
class Solution:
def dfs(self, node, parent, adj, labels, ans):
nodeCounts = [0] * 26
nodeCounts[ord(labels[node]) - ord('a')] = 1
for child in adj[node]:
if child == parent:
continue
childCounts = self.dfs(child, node, adj, labels, ans)
for i in range(26):
nodeCounts[i] += childCounts[i]
ans[node] = nodeCounts[ord(labels[node]) - ord('a')]
return nodeCounts
def countSubTrees(self, n, edges, labels):
from collections import defaultdict
adj = defaultdict(list)
for a, b in edges:
adj[a].append(b)
adj[b].append(a)
ans = [0] * n
self.dfs(0, -1, adj, labels, ans)
return ans
Ставь 👍 и забирай 📚 Базу знаний