All Nodes Distance K in Binary Tree

/*
 本算法答案由上岸科技提供。
 上岸科技是一个专致力于高效培养北美留学生拥有实际面试能力的团体。
 我们采用小班化线上,线下教学让学生更快,更好的学习到想要的知识。
 团队主要由一群怀揣情怀于美国高校毕业的一线IT公司工程师构成。
 我们坚信对于求职者算法并不是全部,合理的技巧加上适当的算法培训能够大大的提升求职成功概率也能大大减少刷题的痛苦。
 正如我们的信仰:我们教的是如何上岸而不仅是算法。
 更多信息请关注官网:https://www.shanganonline.com/
*/
class Solution {
    static class Node {
        int val;
        List<Node> neighbors;
        public Node(int val) {
            this.val = val;
            neighbors = new ArrayList<>();
        }
    }
    
    Node doublyRoot;
    Node doublyTarget;
    Node parent;
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        
        doublyRoot = buildDoublyTree(root, target);
        Queue<Node> queue = new ArrayDeque<>();
        queue.offer(doublyTarget);
        List<Integer> result = new ArrayList<>();
        Set<Integer> visited = new HashSet<>();
        visited.add(doublyTarget.val);
        int level = 0;
        
        while (!queue.isEmpty() && level != K) {
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                Node current = queue.poll();
                for (Node neighbor : current.neighbors) {
                    if (neighbor != null && !visited.contains(neighbor.val)) {
                        visited.add(neighbor.val);
                        queue.offer(neighbor);
                    }
                }
            }
            level++;
        }
        
        while (!queue.isEmpty()) {
            result.add(queue.poll().val);
        }
        return result;
    }
    
    private Node buildDoublyTree(TreeNode root, TreeNode target) {
        if (root == null) {
            return null;
        }
        
        Node current = new Node(root.val);
        current.neighbors.add(parent);
        if (root == target) {
            doublyTarget = current;
        }
        parent = current;
        Node left = buildDoublyTree(root.left, target);
        parent = current;
        Node right = buildDoublyTree(root.right, target);
        current.neighbors.add(left);
        current.neighbors.add(right);
        
        return current;
    }
}

Last updated

Was this helpful?