LeetCode 863 二叉树中所有距离为K的节点

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

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

示例 1：

输入：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" 实际上是树上的结点。

上面的输入仅仅是对这些对象进行了序列化描述。

package com.example.lettcode.tree;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @Class DistanceK
 * @Description 863 二叉树中所有距离为K的节点
 * @Author
 * @Date 2021/7/28
 **/
public class DistanceK {

    static class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        public TreeNode(int val) {
            this.val = val;
            this.left = null;
            this.right = null;
        }
    }

    static Map<TreeNode, TreeNode> treeNodeParentMap = new HashMap<>();
    static List<Integer> ans = new ArrayList<>();


    /**
     * (1) 先找出每个节点的父节点
     * (2) 然后使用bfs，从目标节点出发，找到距离为K的节点
     *
     * @param root
     * @param target
     * @param k
     * @return
     */
    public static List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        if (root == null || target == null) return new ArrayList<>();

        ans = new ArrayList<>();
        treeNodeParentMap = new HashMap<>();

        // 构建节点与其父节点的映射关系
        findParents(root);
        // 从目标节点出发，找到距离为K的节点
        findTargetNode(target, null, 0, k);

        return ans;
    }

    /**
     * 寻找每个节点的父节点
     *
     * @param treeNode
     */
    private static void findParents(TreeNode treeNode) {
        if (treeNode == null) return;
        if (treeNode.left != null) {
            treeNodeParentMap.put(treeNode.left, treeNode);
            findParents(treeNode.left);
        }
        if (treeNode.right != null) {
            treeNodeParentMap.put(treeNode.right, treeNode);
            findParents(treeNode.right);
        }
    }

    /**
     * 从目标节点开始寻找距离为K的节点
     *
     * @param node
     * @param from
     * @param depth
     * @param K
     */
    private static void findTargetNode(TreeNode node, TreeNode from, int depth, int K) {
        if (node == null) return;

        // 该节点与目标节点的距离为K
        if (depth == K) {
            ans.add(node.val);
            return;
        }

        // 分别从左右子树开始寻找距离为K的节点
        if (node.left != from) {
            findTargetNode(node.left, node, depth + 1, K);
        }
        if (node.right != from) {
            findTargetNode(node.right, node, depth + 1, K);
        }
        // 从父节点开始寻找距离为K的节点，就是从父节点的另一子节点开始寻找
        if (treeNodeParentMap.get(node) != from) {
            findTargetNode(treeNodeParentMap.get(node), node, depth + 1, K);
        }
    }
}

测试用例

public static void main(String[] args) {
    // root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
    TreeNode root = new TreeNode(3);
    TreeNode treeNode02 = new TreeNode(5);
    TreeNode treeNode03 = new TreeNode(1);
    root.left = treeNode02;
    root.right = treeNode03;
    TreeNode treeNode04 = new TreeNode(6);
    TreeNode treeNode05 = new TreeNode(2);
    treeNode02.left = treeNode04;
    treeNode02.right = treeNode05;
    TreeNode treeNode06 = new TreeNode(0);
    TreeNode treeNode07 = new TreeNode(8);
    treeNode03.left = treeNode06;
    treeNode03.right = treeNode07;
    TreeNode treeNode08 = new TreeNode(7);
    TreeNode treeNode09 = new TreeNode(4);
    treeNode05.left = treeNode08;
    treeNode05.right = treeNode09;

    TreeNode target = treeNode02;
    int K = 2;
    List<Integer> integerList = DistanceK.distanceK(root, target, K);
    System.out.print("DistanceK demo01 result : [");
    for (Integer integer : integerList) {
        System.out.print(", " + integer);
    }
    System.out.println("]");


    root = new TreeNode(1);
    target = root;
    K = 3;
    integerList = DistanceK.distanceK(root, target, K);
    System.out.print("DistanceK demo02 result : [");
    for (Integer integer : integerList) {
        System.out.print(", " + integer);
    }
    System.out.println("]");
}

运行结果

DistanceK demo01 result : [, 7, 4, 1]
DistanceK demo02 result : []