📃 题目描述

题目链接：863. 二叉树中所有距离为 K 的结点 (opens new window)

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

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

🔔 解题思路

直接从 target 结点处开始走，在每个结点处都有三个可能的前进方向：左孩子方向、右孩子方向，父结点方向，用三分支递归实现。个人认为两个关键点：

1. 由于是二叉链表，所以，无法直接由当前结点走向其父节点，所以用了一个 map 加一次 dfs 遍历来保存所有结点的父节点，这样就可以通过 map 直接跳到父节点了。
2. 为了防止走回头路，需要存储已经访问过的节点（将路过的结点加入 set，若待访问的结点不在 set 中，则访问它，否则跳过）

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

    // value 是 key 的父节点
    private Map<TreeNode, TreeNode> map = new HashMap<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {

        if (root == null) {
            return res;
        }
        
        // 从 root 出发 DFS，记录每个结点的父结点
        findParents(root);
        
        // 存储已访问过的节点
        Set<TreeNode> visited = new HashSet<>();
        // 从 target 出发 DFS，寻找所有深度为 k 的结点，0 d
        findKNodes(target, visited, 0, k);
        
        return res;
    }

    private void findParents(TreeNode cur) {
        if (cur.left != null) {
            map.put(cur.left, cur);
            findParents(cur.left);
        }
        if (cur.right != null) {
            map.put(cur.right, cur);
            findParents(cur.right);
        }
    }

    // 从 node(node 的上一个节点是 from) 出发 DFS，寻找所有深度为 k 的结点
    private void findKNodes(TreeNode node, Set<TreeNode> visited, int depth, int k) {
        if (node == null) {
            return ;
        }

        // 当前节点置为已访问
        visited.add(node);

        if (depth == k) {
            res.add(node.val);
            return ;
        }

        // 处理 node 的邻居节点
        // 左孩子
        if (!visited.contains(node.left)) {
            findKNodes(node.left, visited, depth + 1, k);
        }
        // 右孩子
        if (!visited.contains(node.right)) {
            findKNodes(node.right, visited, depth + 1, k);
        }
        // 父节点
        if (!visited.contains(map.get(node))) {
            findKNodes(map.get(node), visited, depth + 1, k);
        }
    }
}

还可以在空间上优化下，设计一个 from 标志来代替 visited，如下：

class Solution {
    
    private List<Integer> res = new ArrayList<>();

    // value 是 key 的父节点
    private Map<TreeNode, TreeNode> map = new HashMap<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {

        if (root == null) {
            return res;
        }
        
        // 从 root 出发 DFS，记录每个结点的父结点
        findParents(root);
        
        // 从 target 出发 DFS，寻找所有深度为 k 的结点
        findKNodes(target, null, 0, k);
        
        return res;
    }
    
	// dfs
    private void findParents(TreeNode cur) {
        if (cur.left != null) {
            map.put(cur.left, cur);
            findParents(cur.left);
        }
        if (cur.right != null) {
            map.put(cur.right, cur);
            findParents(cur.right);
        }
    }

    // 从 node(node 的上一个节点是 from) 出发 DFS，寻找所有深度为 k 的结点
    private void findKNodes(TreeNode node, TreeNode from, int depth, int k) {
        if (depth == k) {
            res.add(node.val);
            return ;
        }
		
        // 处理 node 的邻居节点
        // 1. node 不是来自 node.left(即 node.left 没有被访问过)，那么就 DFS node.left
        if (node.left != from) {
            findKNodes(node.left, node, depth + 1, k);
        }
        // 2. node 不是来自 node.right(即 node.right 没有被访问过)，那么就 DFS node.right，那么就
        if (node.right != from) {
            findKNodes(node.right, node, depth + 1, k);
        }
        // 3. node 不是来自其父节点(即 node.父节点 没有被访问过)，那么就 DFS 父节点
        if (map.get(node) != from) {
            findKNodes(map.get(node), node, depth + 1, k);
        }

    }
}

💥 复杂度分析

• 空间复杂度：O(N)，记录父节点需要 O(n) 的空间，深度优先搜索需要 O(n) 的栈空间
• 时间复杂度：O(N)，两次 DFS，每次 DFS 时间复杂度 O(v + e)，对应本题就是 O(N)，N 表示节点个数