給出一棵樹的所有邊的集合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