LeetCode863. 二叉树中所有距离为 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

示例 2:

输入: root = [1], target = 1, k = 3
输出: []

提示:

• 节点数在 [1, 500] 范围内
• 0 <= Node.val <= 500
• Node.val 中所有值 不同
• 目标结点 target 是树上的结点。
• 0 <= k <= 1000

解题思路

1. 定义一个递归函数dfs，其输入参数为当前结点node、目标结点target和当前结点到目标结点的距离dist。在函数内部，我们首先检查node是否为空，如果为空则直接返回。接下来，我们检查dist是否等于k，如果等于k，则将当前结点的值添加到结果列表中。否则，我们继续遍历左右子树，更新距离dist。

2. 对于根结点root，我们首先计算它到目标结点target的距离，然后调用递归函数dfs(target, prev, dist)。

3. 在遍历过程中，我们需要记录每个结点的父结点，以便在递归过程中可以通过父结点更新距离。我们可以使用一个字典parent来存储每个结点的父结点。

4. 在递归函数dfs内部，向上遍历时要保证当前节点的左右子节点不是前驱节点，向下遍历时要保证当前节点的父节点不是前驱节点。都符合的话我们可以更新距离dist并继续遍历子树。

5. 当整个二叉树遍历完成后，我们就得到了所有距离为k的结点的值的列表。返回结果列表即可。

注意：在实际实现中，我们可能需要先遍历一次二叉树来构建父结点字典parent。

代码实现

Java广度优先版

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

import java.util.*;

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}

public class Solution {
    Map<TreeNode, TreeNode> parentMap = new HashMap<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        dfs(root, null);
        Queue<TreeNode> queue = new LinkedList<>();
        queue.add(target);
        Set<TreeNode> visited = new HashSet<>();
        visited.add(target);
        int distance = 0;

        while (!queue.isEmpty()) {
            if (distance == K) {
                List<Integer> result = new ArrayList<>();
                for (TreeNode node : queue) {
                    result.add(node.val);
                }
                return result;
            }

            int size = queue.size();
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                if (node.left != null && !visited.contains(node.left)) {
                    visited.add(node.left);
                    queue.add(node.left);
                }
                if (node.right != null && !visited.contains(node.right)) {
                    visited.add(node.right);
                    queue.add(node.right);
                }
                TreeNode parent = parentMap.get(node);
                if (parent != null && !visited.contains(parent)) {
                    visited.add(parent);
                    queue.add(parent);
                }
            }
            distance++;
        }

        return new ArrayList<>();
    }

    private void dfs(TreeNode node, TreeNode parent) {
        if (node == null) return;
        parentMap.put(node, parent);
        dfs(node.left, node);
        dfs(node.right, node);
    }
}

Python深度优先版 

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        # 记录每个结点的父结点
        parent = dict()
        
        # 建立父结点字典
        def find_parent(node):
            if not node:
                return
            if node.left:
                parent[node.left] = node
                find_parent(node.left)
            if node.right:
                parent[node.right] = node
                find_parent(node.right)
        
        find_parent(root)
        
        # 结果列表
        res = []
        
        # 深度优先搜索
        def dfs(node, prev, dist):
            if not node:
                return
            if dist == k:
                res.append(node.val)
                return
            if   node.left != prev:
                dfs(node.left, node, dist + 1)
            if  node.right !=prev:
                dfs(node.right, node, dist + 1)
            if parent.get(node) != prev:
                dfs(parent.get(node), node, dist + 1)
        
        dfs(target, None, 0)
        
        return res