LeetCode 863 All Nodes Distance K in Binary Tree

LeetCode

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer in any order.

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
Output: [7,4,1]
Explanation: The nodes that are a distance 2 from the target node (with value 5) have values 7, 4, and 1.

Example 2:

Input: root = [1], target = 1, k = 3
Output: []

Constraints:

The number of nodes in the tree is in the range [1, 500]. 0 <= Node.val <= 500
All the values Node.val are unique. target is the value of one of the nodes in the tree.
* 0 <= k <= 1000

算法思路：

有三种情况，都容易忽略： 1. 儿子节点 2. 所有父节点路劲上 3. 兄弟节点路径上。而第三种情况要搜另一边的儿子节点（左右不确定）要用visited记录，而且不一定是父亲的兄弟节点，可能爷爷的非父亲的儿子节点。
既然不是单向搜索，不妨转换为图，然后用计算距离BFS模板，只要用map来记录某节点的父亲节点或者增加一个域。BFS中for neighbor in [node.left, node, right, node.parent]

注意事项：

1. root的parent是None，所以从root去赋值parent，而不是从parent root给儿子赋parent
2. Line 13 node.left, node.right, node.parent都可能为None，所以Line14要加not neighbor
3. BFS从target开始而不是root

Python代码：

def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
	if not root:
		return []
	self.dfs(root, None)
	queue, visited, distance_to_tgt,  = collections.deque([target]), set([target]), collections.defaultdict(int)
	distance_to_tgt[target], res = 0, []
	while queue:
		node = queue.popleft()
		if distance_to_tgt[node] == k:
			res.append(node.val)
		if distance_to_tgt[node] > k:
			break
		for neighbor in [node.left, node.right, node.parent]:
			if not neighbor or neighbor in visited:  # remember not neighbor
				continue
			queue.append(neighbor)
			visited.add(neighbor)
			distance_to_tgt[neighbor] = distance_to_tgt[node] + 1
	return res

def dfs(self, root, parent):
	if not root:
		return None
	root.parent = parent
	'''if root.left:
		root.left.parent = root
	if root.right:
		root.right.parent = root'''
	self.dfs(root.left, root)  # remember to pass parent
	self.dfs(root.right, root)

算法分析：

时间复杂度为O(n)，空间复杂度O(n)