[数据结构与算法]Leetcode--Java--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

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */

思路

1. 深度优先遍历 + 哈希表。 由于所给结点没有与父结点的关系，而对于树中某结点需要往父结点的方向遍历。 并且由于结点值是唯一的，因此可以先进行一遍DFS，用哈希表存储每个结点的父结点。 键值为结点的值。
2. 进行第二遍深度优先遍历，从target结点开始，往左子树，右子树，以及父结点方向遍历，当距离等于k时存入结果集。
3. 第二遍遍历时，为了防止重复扫描走过的结点，用一个参数记录来源结点，若来源结点不是即将要 遍历的结点才进行遍历。

代码

class Solution {
    //记录父节点
    Map<Integer, TreeNode> parents = new HashMap<>();

    List<Integer> res = new ArrayList<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        //先进行一遍扫描，记录所有结点的父节点
        getParents(root);
        //以target为根结点查找距离为K的结点，记录上个来源结点防止重复扫描同一个结点
        dfs(target, null, 0, k);       
         return res;
    }

    private void getParents(TreeNode node) {
        if (node.left != null) {
            parents.put(node.left.val, node);
            getParents(node.left);
        }
        if (node.right != null) {
            parents.put(node.right.val, node);
            getParents(node.right);
        }
    }

    private void dfs(TreeNode node, TreeNode from, int depth, int k) {
        //为空就返回
        if (node == null) {
            return;
        }
        if (depth == k) {
            res.add(node.val);
            return;
        }
        //扫描左子树
        if (node.left != from) {
              dfs(node.left, node, depth + 1, k);
        }
        //扫描右子树
        if (node.right != from) {
            dfs(node.right, node, depth + 1, k);
        }
        //往根结点扫描
        if (parents.get(node.val) != from) {
            dfs(parents.get(node.val), node, depth + 1, k);
        }
    }


}