All Nodes Distance K in Binary Tree

I had come across this question in an interview some years back and recently in Leetcode. I recently got interested in solving problems but this time I thought to do it differently. I think there is an efficient technique to reach any node from a given target with any constraint.

To solve this problem we should track node parents and the path from the root to target, Go to each parent and find all Node distance ( (k -distance) from target ).

Instead of storing parents and path, Why don’t we take advantage of the path having all parent-child nodes?

But then when finding l ( k — distance from target) distance from a parent, we may count the l-1 nodes from a child as well which are not correct nodes.

To solve this problem, we can just delink the child, count l distance from the parent, link again and move to parent’s parent. Do this process till we reach to root.

Update: We actually don’t need to store path from root to target but with DFS we can figure out if the current Node is in the path and also the child. We just have to how much distance is between the target to current and find the All Nodes k-d distance Nodes.

First version code is at the bottom. Below one is a more efficient version but uses the same algorithm.

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public int distanceFromTarget = 0;
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        List<Integer> ans = new ArrayList<>();
        distanceKFromRoot(target, K, ans);
        dfsutil(root, target, K, ans);
        return ans;
    }
    
    public TreeNode dfsutil(TreeNode root, TreeNode target, int K, List<Integer> ans) {
        if( root == null) return null;
        if(root == target) {
           distanceFromTarget++; 
           return root;
        }
        TreeNode leftChildInPath = dfsutil(root.left, target, K, ans);
        if(leftChildInPath != null) {
            root.left = null;
            if(K >= distanceFromTarget) {
                distanceKFromRoot(root, K-distanceFromTarget, ans); 
            }
            root.left = leftChildInPath;
            distanceFromTarget++;
            return root;
        }
        TreeNode rightChildInPath = dfsutil(root.right, target, K, ans);
        if(rightChildInPath != null) {
            root.right = null;
            if(K >= distanceFromTarget) {
                distanceKFromRoot(root, K-distanceFromTarget, ans); 
            }
            root.right = rightChildInPath;
            distanceFromTarget++;
            return root;
        }
        return null;
    }
    
    public void distanceKFromRoot(TreeNode root, int k, List<Integer> ans) {
        if(root == null) return;
        if(k == 0) {
            ans.add(root.val);
            return;
        }
        distanceKFromRoot(root.left, k-1, ans);
        distanceKFromRoot(root.right, k-1, ans);
    }
}

First Version:

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        List<Integer> ans = new ArrayList<>();
        distanceKFromRoot(target, K, ans);
        List<TreeNode> path = new ArrayList<>();
        dfsutil(root, target, K, path, ans);
        return ans;
    }
    
    public boolean dfsutil(TreeNode root, TreeNode target, int K, List<TreeNode> path, List<Integer> ans) {
        if( root == null) return false;
        if(root == target) {
            int count = K - 1;
            TreeNode child = target;
            for( int i = path.size() -1; i >= 0 && count >= 0; i--) {
                TreeNode current = path.get(i);
                if(current.left == child) {
                    current.left = null;
                    distanceKFromRoot(current, count, ans);   
                    current.left = child;
                } else if( current.right == child) {
                    current.right = null;
                    distanceKFromRoot(current, count, ans);   
                    current.right = child;
                }
                child = current;
                count--;
            }
            return true;
        }
        path.add(root);
        boolean ret = dfsutil(root.left, target, K, path, ans);
        path.remove(path.size()-1);
        if(ret) return true;
        path.add(root);
        ret = dfsutil(root.right, target, K, path, ans);
        path.remove(path.size()-1);
        return ret;
    }
    
    public void distanceKFromRoot(TreeNode root, int k, List<Integer> ans) {
        if(root == null) return;
        if(k == 0) {
            ans.add(root.val);
            return;
        }
        distanceKFromRoot(root.left, k-1, ans);
        distanceKFromRoot(root.right, k-1, ans);
    }
}

Hope you find this technique useful.