求二叉树中和某个target节点距离为k的所有节点值,先BFS建图(python的dict存图),然后在图上BFS求解
1 #863
2 #Runtime: 25 ms (Beats 49.62%)
3 #Memory: 13.6 MB (Beats 74.5%)
4
5 # Definition for a binary tree node.
6 # class TreeNode(object):
7 # def __init__(self, x):
8 # self.val = x
9 # self.left = None
10 # self.right = None
11
12 class Solution(object):
13 def distanceK(self, root, target, k):
14 """
15 :type root: TreeNode
16 :type target: TreeNode
17 :type k: int
18 :rtype: List[int]
19 """
20 graph = {}
21 dep_target = 0
22 q = deque([[root, -1]])
23 while q:
24 node, f = q.popleft()
25 if node.val not in graph:
26 if f != -1:
27 graph[node.val] = [f]
28 graph[f].append(node.val)
29 else:
30 graph[node.val] = []
31 else:
32 graph[node.val].append(f)
33 graph[f].append(node.val)
34 if node.left:
35 q.append([node.left, node.val])
36 if node.right:
37 q.append([node.right, node.val])
38 q = deque([target.val])
39 dis = 0
40 ans = []
41 vis = set([target.val])
42 while q and dis < k:
43 sz = len(q)
44 while sz:
45 sz -= 1
46 node = q.popleft()
47 for i in graph[node]:
48 if i not in vis:
49 q.append(i)
50 vis.add(i)
51 dis += 1
52 if dis == k:
53 ans = list(q)
54 return ans