LeetCode 863
这个题目应该还是挺直接的,因为是要找K distance的所有node,只需要用BFS做K层查找就可以了。
唯一的一点是怎么找回parent节点,这个我们可以使用一次树的遍历建立和parent直接的hash table。
def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:
if K == 0: return [target.val]
parents = collections.defaultdict(TreeNode)
parents[root.val] = None
def findParents(node: TreeNode):
if node == None: return None
if node.left != None:
parents[node.left.val] = node
findParents(node.left)
if node.right != None:
parents[node.right.val] = node
findParents(node.right)
def appendNode(q, node, visited, distance):
if node != None and node not in visited:
q.append((node, distance))
findParents(root)
q = [(target,0)]
#BFS
result = []
visited = set()
while q:
currentNode, distance = q.pop(0)
visited.add(currentNode)
if distance == K:
result.append(currentNode.val)
else:
appendNode(q, parents[currentNode.val], visited, distance + 1)
appendNode(q, currentNode.left, visited, distance+1)
appendNode(q, currentNode.right, visited, distance+1)
return result