力扣 863. 二叉树中所有距离为 K 的结点

题目

给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 K 。

返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。

示例

输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
输出：[7,4,1]
解释：
所求结点为与目标结点（值为 5）距离为 2 的结点，
值分别为 7，4，以及 1

注意，输入的 “root” 和 “target” 实际上是树上的结点。
上面的输入仅仅是对这些对象进行了序列化描述。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

方法1：图的思想

Java实现

class Solution {
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        //哈希表存 子结点-父结点
        Map<TreeNode, TreeNode> map = new HashMap<>();
        Queue<TreeNode> q = new LinkedList<>();
        q.offer(root);

        while (!q.isEmpty()) {
            int sz = q.size();
            for (int i = 0; i < sz; i++) {
                TreeNode cur = q.poll();
                if (cur.left != null) {
                    map.put(cur.left, cur);
                    q.offer(cur.left);
                }
                if (cur.right != null) {
                    map.put(cur.right, cur);
                    q.offer(cur.right);
                }
            }
        }

        //哈希表+树结合起来，实现BFS遍历
        int step = 0;
        List<Integer> res = new ArrayList<>();

        if (step == k) {
            res.add(target.val);
            return res;
        }

        //标识已经遍历过的结点，防止重复遍历
        Set<TreeNode> set = new HashSet<>();
        q.offer(target);
        set.add(target);

        while (!q.isEmpty()) {
            int sz = q.size();
            for (int i = 0; i < sz; i++) {
                TreeNode cur = q.poll();

                //该结点的父结点入队列
                if (map.containsKey(cur) && !set.contains(map.get(cur))) {
                    q.offer(map.get(cur));
                    set.add(map.get(cur));
                }
                //该结点的左孩子、右孩子入队列
                if (cur.left != null && !set.contains(cur.left)) {
                    q.offer(cur.left);
                    set.add(cur.left);
                }
                if (cur.right != null && !set.contains(cur.right)) {
                    q.offer(cur.right);
                    set.add(cur.right);
                }
            }
            step++;
            if (step == k) {
                while (!q.isEmpty()) {
                    res.add(q.poll().val);
                }
            }
        }

        return res;
    }
}