863. All Nodes Distance K in Binary Tree

/**
 * 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> res = new ArrayList<>();
        if (root == null || K < 0) {
            return res;
        }
        Map<TreeNode, List<TreeNode>> map = new HashMap<>();
        buildGraph(root, null, map);
        
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(target);
        // int step = 0;
        Set<TreeNode> visited = new HashSet<>();
        visited.add(target);
        while (!queue.isEmpty()) {
            int size = queue.size();
            if (K == 0) {
                for (int i = 0; i < size; i++) {
                    res.add(queue.poll().val);
                }
                return res;
            }
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                for (TreeNode next : map.get(node)) {
                    if (visited.contains(next)) {
                        continue;
                    }
                    visited.add(next);
                    queue.offer(next);
                }
            }
            K--;
        }
        return res;
    }
    
    private void buildGraph(TreeNode node, TreeNode parent, Map<TreeNode, List<TreeNode>> map) {
        if (node == null) {
            return;
        }
        if (!map.containsKey(node)) {
            map.put(node, new ArrayList<>());
            if (parent != null) {
                map.get(node).add(parent);
                map.get(parent).add(node);
            }
            buildGraph(node.left, node, map);
            buildGraph(node.right, node, map);
        }
    }
}

/**
 * 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> res = new ArrayList<>();
        if (root == null) {
            return res;
        }
        Map<TreeNode, TreeNode> map = new HashMap<>();
        findParent(root, null, map);
        Set<TreeNode> visited = new HashSet<>();
        helper(target, K, map, visited, res);
        return res;
    }
    
    private void helper(TreeNode node, int K, Map<TreeNode, TreeNode> map, Set<TreeNode> visited, List<Integer> res) {
        if (visited.contains(node)) {
            return;
        }
        visited.add(node);
        if (K == 0) {
            res.add(node.val);
            return;
        }
        if (node.left != null) {
            helper(node.left, K - 1, map, visited, res);
        }
        if (node.right != null) {
            helper(node.right, K - 1, map, visited, res);
        }
        if (map.containsKey(node)) {
            helper(map.get(node), K - 1, map, visited, res);
        }
    }
    
    private void findParent(TreeNode node, TreeNode parent, Map<TreeNode, TreeNode> map) {
        if (node == null) {
            return;
        }
        if (node.left != null) {
            map.put(node.left, node);
        }
        if (node.right != null) {
            map.put(node.right, node);
        }
        findParent(node.left, node, map);
        findParent(node.right, node, map);
    }
}

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    private List<Integer> res = new ArrayList<>();
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        dis(root, target, K);
        return res;
    }
    
    // distance from root to target; returns -1 if target is not in the tree
    private int dis(TreeNode root, TreeNode target, int K) {
        if (root == null) {
            return -1;
        }
        if (root == target) {
            collect(target, K);
            return 0;
        }
        int left = dis(root.left, target, K);
        int right = dis(root.right, target, K);
        
        // Target in the left subtree
        if (left >= 0) {
            if (left == K - 1) {
                res.add(root.val); 
            }
            collect(root.right, K - left - 2);
            return left + 1;
        }
        
        // Target in the right subtree
        if (right >= 0) {
            if (right == K - 1) {
                res.add(root.val);
            }
            collect(root.left, K - right - 2);
            return right + 1;
        }
        return -1;
    }
    
    // Collect nodes that are d steps from root
    private void collect(TreeNode root, int d) {
        if (root == null || d < 0) {
            return;
        }
        if (d == 0) {
            res.add(root.val);
        }
        collect(root.left, d - 1);
        collect(root.right, d - 1);
    }
}