【DFS】LeetCode 863. 二叉树中所有距离为 K 的结点

题目链接

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

思路

因为在二叉树中，每个结点至多连接三条边，并且其中两条已经通过左右子结点的形式给出，所以我们只需要记录每个结点的父结点即可。

使用 Map 记录每个结点的父结点，再直接以 target 作为源点开始 DFS，记录与 target 距离为 k 的结点。

在传递 DFS 参数的过程中我们需要传入变量 from 记录当前结点是从哪个结点遍历而来的，防止重复遍历。

代码

class Solution {
    Map<Integer, TreeNode> parents = new HashMap<>();
    List<Integer> result = new ArrayList<>();

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        buildGraph(root);
        dfs(target, null, 0, k);

        return result;
    }

    private void dfs(TreeNode node, TreeNode from, int depth, int k) {
        if(node == null){
            return;
        }
        if(depth == k){
            result.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);
        }
    }

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