863. All Nodes Distance K in Binary Tree **

JIAKAOBO

LeetCode

感谢赞助!

• ㊗️
• 大家
• offer
• 多多！

Problem

We are given a binary tree (with root node root), a target node, and an integer value k.

Return a list of the values of all nodes that have a distance k from the target node. The answer can be returned in any order.

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2

Output: [7,4,1]

Explanation:
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.

Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.

Constraints:

• The number of nodes in the tree is in the range [1, 500].
• 0 <= Node.val <= 500
• All the values Node.val are unique.
• target is the value of one of the nodes in the tree.
• 0 <= k <= 1000

Code

class Solution {
    HashMap<Integer, List<Integer>> map = new HashMap<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        // 建图
        buildGraph(null, root);
        List<Integer> res = new ArrayList<>();

        HashSet<Integer> visited = new HashSet<>();
        Queue<Integer> queue = new LinkedList<>();

        queue.offer(target.val);
        visited.add(target.val);

        int k = 0;
        while (!queue.isEmpty() && k <= K) {
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                int curr = queue.poll();
                if (k == K) {
                    res.add(curr);
                }
                if (!map.containsKey(curr)) continue;
                for (int next : map.get(curr)) {
                    if (visited.contains(next)) continue;
                    queue.offer(next);
                    visited.add(next);
                }
            }

            k++;
        }

        return res;
    }

    // 如何把tree变成graph
    private void buildGraph(TreeNode parent, TreeNode child) {
        if (parent != null) {
            if (!map.containsKey(parent.val)) {
                map.put(parent.val, new ArrayList<>());
            }

            map.get(parent.val).add(child.val);

            if (!map.containsKey(child.val)) {
                map.put(child.val, new ArrayList<>());
            }

            map.get(child.val).add(parent.val);
        }

        if (child.left != null) buildGraph(child, child.left);
        if (child.right != null) buildGraph(child, child.right);
    }
}

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 * int val;
 * TreeNode left;
 * TreeNode right;
 * TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    List<Integer> res = new ArrayList<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        helper(root, target, K);
        return res;
    }

    private int helper(TreeNode root, TreeNode target, int K) {
        if (root == null) return -1;
        // 找到了target
        // 那么要寻找距离target K的点
        if (root == target) {
            collect(target, K);
            return 0;
        }

        int l = helper(root.left, target, K);
        int r = helper(root.right, target, K);

        if (l >= 0) {
            // l到target的距离是1
            // 那么当前节点就是解
            if (l == K - 1) {
                res.add(root.val);
            }

            collect(root.right, K - l - 2);
            return l + 1;
        }

        if (r >= 0) {
            if (r == K - 1) {
                res.add(root.val);
            }

            collect(root.left, K - r - 2);
            return r + 1;
        }

        return -1;
    }

    // 找到root的子节点,距离为d
    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);
    }
}