Problem Statement

All Nodes Distance K in Binary Tree - LeetCode Pattern:


Solution

Map<Node, Integer> map = new HashMap<>();
List<Integer> res;
 
public List<Integer> distanceK (Node root, Node target, int K){
	res = new LinkedList<>();
	
	find(root, target);            // find distances for all in target path
	dfs(root, map.get(root), K);   // dfs tree to find all k-distant nodes to target
	
	return res;
}
 
private void dfs (Node root, int dist, int K) { // dist -> distance from Target
	if(root == null) return;
	
	// override dist if in target path
	dist = map.getOrDefault(root, dist);
	
	// if dist == k ans found, append result
	if(dist == K) res.add(root.val);
 
	// recurse for left right
	dfs(root.left, dist+1, K);  // inc dist by 1
	dfs(root.right, dist+1, K); // inc dist by 1
}
 
// find distances for all in target path from root
private int find(Node root, Node target) {
	// BASE CASES
	if(root == null) return -1;
	if(root == target) {        // target found
		map.put(root, 0);
		return 0;
	}
	// recurse for left
	int leftDist = find(root.left, target);
	if (leftDist > -1){             // leftDist to target found
		map.put(root, leftDist+1);  // map -> root, leftDist
		return leftDist+1;
	}
	// recurse for right
	int rightDist = find(root.right, target);
	if(rightDist > -1) {            // rightDist to target found
		map.put(root, rightDist+1); // map -> root, rightDist
		return rightDist+1;
	}
	return -1;  // target not found in subtree
}

Notes

As we know, if the distance from a node to target node is k, the distance from its child to the target node is k + 1 unless the child node is closer to the target node which means the target node is in it’s subtree.

To avoid this situation, we need to travel the tree first to find the path from root to target, to:

  • store the value of distance in hashamp from the all nodes in that path to target

Then we can easily use dfs to travel the whole tree. Every time when we meet a treenode which has already been stored in map, use the stored value in hashmap instead of plus 1. and the distance to it’s child node is ++dist

Still dont understand :Nodes at Distance K in Binary Tree | C++ Placement Course | Lecture 27.14 - YouTube