863. All Nodes Distance K in Binary Tree

time: buildGraph: O(n) + bfs: O(n) = O(n)

space: O(n)

/**
 * 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, List<TreeNode>> map = new HashMap<>();
        buildGraph(root, null, map);
        
        if (!map.containsKey(target)) { // no target
            return res;
        }
        
        return bfs(map, target, k, res);
        
    }
    // do BFS
    private List<Integer> bfs(Map<TreeNode, List<TreeNode>> map, TreeNode target, int k, List<Integer> res) {
        Set<TreeNode> visited = new HashSet<>();
        Queue<TreeNode> queue = new ArrayDeque<>();
        visited.add(target);
        queue.offer(target);
        
        while (!queue.isEmpty()) {
            int size = queue.size();
            
            for (int i = 0; i < size; i++) { 
                if (k == 0) {
                    while (!queue.isEmpty()) {
                        res.add(queue.poll().val);
                    }
                    return res;
                }
                
                TreeNode curr = queue.poll();
                for (TreeNode node : map.get(curr)) {
                    if (visited.contains(node)) {
                        continue;
                    }
                    visited.add(node);
                    queue.offer(node);
                }
            }
            k--;
        }
        return res;
    }
    
    private void buildGraph(TreeNode node, TreeNode parent, Map<TreeNode, List<TreeNode>> map) {
        if (node == null) {
            return;
        }
        map.computeIfAbsent(node, n -> new ArrayList<>());
        if (parent != null) { // build graph with parent and children
            map.get(node).add(parent);
            map.get(parent).add(node);
        }
        buildGraph(node.left, node, map);
        buildGraph(node.right, node, map);
    }
}

new version

之前是建map 時把 parent , left, right 都建立

這裡只建立 parent, 其他 left, right 當跑BFS在加入(比較省空間!

```java
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    class Step {
        TreeNode node;
        int level;
        Step(TreeNode node, int level) {
            this.node = node;
            this.level = level;
        }
    }
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        List<Integer> result = new ArrayList<>();

        Map<Integer, TreeNode> parent = new HashMap<>();
        setParent(parent, root, null);

        Queue<Step> queue = new LinkedList<>();
        queue.offer(new Step(target, 0));
        Set<Integer> visited = new HashSet<>();

        while (!queue.isEmpty()) {
            Step cur = queue.poll();

            if (cur.node == null || visited.contains(cur.node.val)) {
                continue;
            } 

            int curVal = cur.node.val;
            if (cur.level == k) {
                result.add(curVal);
                continue;
            }

            visited.add(curVal);
            TreeNode p = parent.get(curVal);
            queue.offer(new Step(p, cur.level + 1));
            queue.offer(new Step(cur.node.left, cur.level + 1));
            queue.offer(new Step(cur.node.right, cur.level + 1));
        }
        return result;
    }
    
    private void setParent(Map<Integer, TreeNode> map, TreeNode root, TreeNode parent) {
        if (root == null) {
            return;
        }
        map.put(root.val, parent);
        setParent(map, root.left, root);
        setParent(map, root.right, root);
    }
}
/*
using dfs to build parent map

then BFS start from target to find level is k

O(n + n)
O(n)
*/
```

Last updated