863. All Nodes Distance K in Binary Tree #

Problem #

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned 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.

Note:

1. The given tree is non-empty.
2. Each node in the tree has unique values 0 <= node.val <= 500.
3. The target node is a node in the tree.
4. 0 <= K <= 1000.

Problem Summary #

Given a binary tree (with root node root), a target node target, and an integer value K. Return a list of the values of all nodes that have a distance K from the target node target. The answer can be returned in any order.

Notes:

• The given tree is non-empty.
• Each node in the tree has a unique value 0 <= node.val <= 500 .
• The target node target is a node in the tree.
• 0 <= K <= 1000.

Solution Ideas #

• Given a tree, a target node target, and a distance K, find all nodes whose distance from the target node target is K.
• This problem is solved using DFS. First find the distance from the current node to the target node. If target is found in the left subtree and the distance from the current node is > 0, then you also need to search its right subtree for the remaining distance. If target is found in the right subtree, the same logic applies in reverse. If the current node is the target node, then you can directly record this node. Otherwise, each time a node is traversed, the distance is decreased by one.

Code #

func distanceK(root *TreeNode, target *TreeNode, K int) []int {
	visit := []int{}
	findDistanceK(root, target, K, &visit)
	return visit
}

func findDistanceK(root, target *TreeNode, K int, visit *[]int) int {
	if root == nil {
		return -1
	}
	if root == target {
		findChild(root, K, visit)
		return K - 1
	}
	leftDistance := findDistanceK(root.Left, target, K, visit)
	if leftDistance == 0 {
		findChild(root, leftDistance, visit)
	}
	if leftDistance > 0 {
		findChild(root.Right, leftDistance-1, visit)
		return leftDistance - 1
	}
	rightDistance := findDistanceK(root.Right, target, K, visit)
	if rightDistance == 0 {
		findChild(root, rightDistance, visit)
	}
	if rightDistance > 0 {
		findChild(root.Left, rightDistance-1, visit)
		return rightDistance - 1
	}
	return -1
}

func findChild(root *TreeNode, K int, visit *[]int) {
	if root == nil {
		return
	}
	if K == 0 {
		*visit = append(*visit, root.Val)
	} else {
		findChild(root.Left, K-1, visit)
		findChild(root.Right, K-1, visit)
	}
}