给出一棵树的所有边的集合edges、节点数、每个节点的label(英文小写字母)、0为根节点、问每个节点的子树中有多少个节点和当前节点的label值一样(包含当前节点)
先根据edges建树(用dict存每个节点的邻接节点),然后从根开始DFS,用Counter()记录每个节点的子树中各个label的节点数的计数
1 #1519
2 #Runtime: 3033 ms (Beats 71.43%)
3 #Memory: 182.4 MB (Beats 57.14%)
4
5 class Solution(object):
6 def countSubTrees(self, n, edges, labels):
7 """
8 :type n: int
9 :type edges: List[List[int]]
10 :type labels: str
11 :rtype: List[int]
12 """
13 ans = [0] * n
14 node = defaultdict(list)
15 for x, y in edges:
16 node[x].append(y)
17 node[y].append(x)
18
19 def DFS(r, p):
20 cnt = Counter()
21 for i in node[r]:
22 if i != p:
23 cnt += DFS(i, r)
24 cnt[labels[r]] += 1
25 ans[r] = cnt[labels[r]]
26 return cnt
27
28 DFS(0, -1)
29 return ans
30