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

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

提示：

给定的树是非空的，且最多有 K 个结点。

树上的每个结点都具有唯一的值 0 <= node.val <= 500 。

目标结点 target 是树上的结点。

0 <= K <= 1000.

来源：力扣（LeetCode）

链接：https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree

一

这道题给了我们一棵二叉树，一个目标结点 target，还有一个整数K，让返回所有跟目标结点 target 相距K的结点。我们知道在子树中寻找距离为K的结点很容易，因为只需要一层一层的向下遍历即可，难点就在于符合题意的结点有可能是祖先结点，或者是在旁边的兄弟子树中，这就比较麻烦了，因为二叉树只有从父结点到子结点的路径，反过来就不行。既然没有，我们就手动创建这样的反向连接即可，这样树的遍历问题就转为了图的遍历（其实树也是一种特殊的图）。建立反向连接就是用一个 HashMap 来来建立每个结点和其父结点之间的映射，使用先序遍历建立好所有的反向连接，然后再开始查找和目标结点距离K的所有结点，这里需要一个 HashSet 来记录所有已经访问过了的结点。

在递归函数中，首先判断当前结点是否已经访问过，是的话直接返回，否则就加入到 visited 中。再判断此时K是否为0，是的话说明当前结点已经是距离目标结点为K的点了，将其加入结果 res 中，然后直接返回。否则分别对当前结点的左右子结点调用递归函数，注意此时带入 K-1，这两步是对子树进行查找。之前说了，还得对父结点，以及兄弟子树进行查找，这是就体现出建立的反向连接 HashMap 的作用了，若当前结点的父结点存在，我们也要对其父结点调用递归函数，并同样带入 K-1，这样就能正确的找到所有满足题意的点了，参见代码如下：

解法一：

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        if (!root) return {};
        vector<int> res;
        unordered_map<TreeNode*, TreeNode*> parent;
        unordered_set<TreeNode*> visited;
        findParent(root, parent);
        helper(target, K, parent, visited, res);
        return res;
    }
    void findParent(TreeNode* node, unordered_map<TreeNode*, TreeNode*>& parent) {
        if (!node) return;
        if (node->left) parent[node->left] = node;
        if (node->right) parent[node->right] = node;
        findParent(node->left, parent);
        findParent(node->right, parent);
    }
    void helper(TreeNode* node, int K, unordered_map<TreeNode*, TreeNode*>& parent, unordered_set<TreeNode*>& visited, vector<int>& res) {
        if (visited.count(node)) return;
        visited.insert(node);
        if (K == 0) {res.push_back(node->val); return;}
        if (node->left) helper(node->left, K - 1, parent, visited, res);
        if (node->right) helper(node->right, K - 1, parent, visited, res);
        if (parent[node]) helper(parent[node], K - 1, parent, visited, res);
    }
};

既然是图的遍历，那就也可以使用 BFS 来做，为了方便起见，我们直接建立一个邻接链表，即每个结点最多有三个跟其相连的结点，左右子结点和父结点，使用一个 HashMap 来建立每个结点和其相邻的结点数组之间的映射，这样就几乎完全将其当作图来对待了，建立好邻接链表之后，原来的树的结构都不需要用了。既然是 BFS 进行层序遍历，就要使用队列 queue，还要一个 HashSet 来记录访问过的结点。在 while 循环中，若K为0了，说明当前这层的结点都是符合题意的，就把当前队列中所有的结点加入结果 res，并返回即可。否则就进行层序遍历，取出当前层的每个结点，并在邻接链表中找到和其相邻的结点，若没有访问过，就加入 visited 和 queue 中即可。记得每层遍历完成之后，K要自减1，参见代码如下：

解法二：

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        if (!root) return {};
        vector<int> res;
        unordered_map<TreeNode*, vector<TreeNode*>> m;
        queue<TreeNode*> q{{target}};
        unordered_set<TreeNode*> visited{{target}};
        findParent(root, NULL, m);
        while (!q.empty()) {
            if (K == 0) {
                for (int i = q.size(); i > 0; --i) {
                    res.push_back(q.front()->val); q.pop();
                }
                return res;
            }
            for (int i = q.size(); i > 0; --i) {
                TreeNode *t = q.front(); q.pop();
                for (TreeNode *node : m[t]) {
                    if (visited.count(node)) continue;
                    visited.insert(node);
                    q.push(node);
                }
            }
            --K;
        }
        return res;
    }
    void findParent(TreeNode* node, TreeNode* pre, unordered_map<TreeNode*, vector<TreeNode*>>& m) {
        if (!node) return;
        if (m.count(node)) return;
        if (pre) {
            m[node].push_back(pre);
            m[pre].push_back(node);
        }
        findParent(node->left, node, m);
        findParent(node->right, node, m);
    }
};

其实这道题也可以不用 HashMap，不建立邻接链表，直接在递归中完成所有的需求，真正体现了递归的博大精深。在进行递归之前，我们要先判断一个 corner case，那就是当 K==0 时，此时要返回的就是目标结点值本身，可以直接返回。否则就要进行递归了。这里的递归函数跟之前的有所不同，是需要返回值的，这个返回值表示的含义比较复杂，若为0，表示当前结点为空或者当前结点就是距离目标结点为K的点，此时返回值为0，是为了进行剪枝，使得不用对其左右子结点再次进行递归。当目标结点正好是当前结点的时候，递归函数返回值为1，其他的返回值为当前结点离目标结点的距离加1。还需要一个参数 dist，其含义为离目标结点的距离，注意和递归的返回值区别，这里不用加1，且其为0时候不是为了剪枝，而是真不知道离目标结点的距离。

在递归函数中，首先判断若当前结点为空，则直接返回0。然后判断 dist 是否为k，是的话，说目标结点距离当前结点的距离为K，是符合题意的，需要加入结果 res 中，并返回0，注意这里返回0是为了剪枝。否则判断，若当前结点正好就是目标结点，或者已经遍历过了目标结点（表现为 dist 大于0），那么对左右子树分别调用递归函数，并将返回值分别存入 left 和 right 两个变量中。注意此时应带入 dist+1，因为是先序遍历，若目标结点之前被遍历到了，那么说明目标结点肯定不在当前结点的子树中，当前要往子树遍历的话，肯定离目标结点又远了一些，需要加1。若当前结点不是目标结点，也还没见到目标结点时，同样也需要对左右子结点调用递归函数，但此时 dist 不加1，因为不确定目标结点的位置。若 left 或者 right 值等于K，则说明目标结点在子树中，且距离当前结点为K（为啥呢？因为目标结点本身是返回1，所以当左右子结点返回K时，和当前结点距离是K）。接下来判断，若当前结点是目标结点，直接返回1，这个前面解释过了。然后再看 left 和 right 的值是否大于0，若 left 值大于0，说明目标结点在左子树中，我们此时就要对右子结点再调用一次递归，并且 dist 带入 left+1，同理，若 right 值大于0，说明目标结点在右子树中，我们此时就要对左子结点再调用一次递归，并且 dist 带入 right+1。这两步很重要，是之所以能不建立邻接链表的关键所在。若 left 大于0，则返回 left+1，若 right 大于0，则返回 right+1，否则就返回0，参见代码如下：

解法三：

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        if (K == 0) return {target->val};
        vector<int> res;
        helper(root, target, K, 0, res);
        return res;
    }
    int helper(TreeNode* node, TreeNode* target, int k, int dist, vector<int>& res) {
        if (!node) return 0;
        if (dist == k) {res.push_back(node->val); return 0;}
        int left = 0, right = 0;
        if (node->val == target->val || dist > 0) {
            left = helper(node->left, target, k, dist + 1, res);
            right = helper(node->right, target, k, dist + 1, res);
        } else {
            left = helper(node->left, target, k, dist, res);
            right = helper(node->right, target, k, dist, res);
        }
        if (left == k || right == k) {res.push_back(node->val); return 0;}
        if (node->val == target->val) return 1;
        if (left > 0) helper(node->right, target, k, left + 1, res);
        if (right > 0) helper(node->left, target, k, right + 1, res);
        if (left > 0 || right > 0) return left > 0 ? left + 1 : right + 1;
        return 0;
    }
};

二

方法一： 深度优先搜索

思路

如果节点有指向父节点的引用，也就知道了距离该节点 1 距离的所有节点。之后就可以从 target 节点开始进行深度优先搜索了。

算法

对所有节点添加一个指向父节点的引用，之后做深度优先搜索，找到所有距离 target 节点 K 距离的节点。

class Solution {
    Map<TreeNode, TreeNode> parent;
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        parent = new HashMap();
        dfs(root, null);

        Queue<TreeNode> queue = new LinkedList();
        queue.add(null);
        queue.add(target);

        Set<TreeNode> seen = new HashSet();
        seen.add(target);
        seen.add(null);

        int dist = 0;
        while (!queue.isEmpty()) {
            TreeNode node = queue.poll();
            if (node == null) {
                if (dist == K) {
                    List<Integer> ans = new ArrayList();
                    for (TreeNode n: queue)
                        ans.add(n.val);
                    return ans;
                }
                queue.offer(null);
                dist++;
            } else {
                if (!seen.contains(node.left)) {
                    seen.add(node.left);
                    queue.offer(node.left);
                }
                if (!seen.contains(node.right)) {
                    seen.add(node.right);
                    queue.offer(node.right);
                }
                TreeNode par = parent.get(node);
                if (!seen.contains(par)) {
                    seen.add(par);
                    queue.offer(par);
                }
            }
        }

        return new ArrayList<Integer>();
    }

    public void dfs(TreeNode node, TreeNode par) {
        if (node != null) {
            parent.put(node, par);
            dfs(node.left, node);
            dfs(node.right, node);
        }
    }
}

Python

class Solution(object):
    def distanceK(self, root, target, K):
        def dfs(node, par = None):
            if node:
                node.par = par
                dfs(node.left, node)
                dfs(node.right, node)

        dfs(root)

        queue = collections.deque([(target, 0)])
        seen = {target}
        while queue:
            if queue[0][1] == K:
                return [node.val for node, d in queue]
            node, d = queue.popleft()
            for nei in (node.left, node.right, node.par):
                if nei and nei not in seen:
                    seen.add(nei)
                    queue.append((nei, d+1))

        return []

复杂度分析

时间复杂度： O(N)，其中 NN 是树中节点个数。

空间复杂度： O(N)。

方法二： 计算节点之间距离

思路

如果 target 节点在 root 节点的左子树中，且 target 节点深度为 3，那所有 root 节点右子树中深度为 K - 3 的节点到 target 的距离就都是 K。

算法

深度优先遍历所有节点。定义方法 dfs(node)，这个函数会返回 node 到 target 的距离。在 dfs(node) 中处理下面四种情况：

如果 node == target，把子树中距离 target 节点距离为 K 的所有节点加入答案。

如果 target 在 node 左子树中，假设 target 距离 node 的距离为 L+1，找出右子树中距离 target 节点 K - L - 1 距离的所有节点加入答案。

如果 target 在 node 右子树中，跟在左子树中一样的处理方法。

如果 target 不在节点的子树中，不用处理。

实现的算法中，还会用到一个辅助方法 subtree_add(node, dist)，这个方法会将子树中距离节点 node K - dist 距离的节点加入答案。

Java

class Solution {
    List<Integer> ans;
    TreeNode target;
    int K;
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        ans = new LinkedList();
        this.target = target;
        this.K = K;
        dfs(root);
        return ans;
    }

    // Return vertex distance from node to target if exists, else -1
    // Vertex distance: the number of vertices on the path from node to target
    public int dfs(TreeNode node) {
        if (node == null)
            return -1;
        else if (node == target) {
            subtree_add(node, 0);
            return 1;
        } else {
            int L = dfs(node.left), R = dfs(node.right);
            if (L != -1) {
                if (L == K) ans.add(node.val);
                subtree_add(node.right, L + 1);
                return L + 1;
            } else if (R != -1) {
                if (R == K) ans.add(node.val);
                subtree_add(node.left, R + 1);
                return R + 1;
            } else {
                return -1;
            }
        }
    }

    // Add all nodes 'K - dist' from the node to answer.
    public void subtree_add(TreeNode node, int dist) {
        if (node == null) return;
        if (dist == K)
            ans.add(node.val);
        else {
            subtree_add(node.left, dist + 1);
            subtree_add(node.right, dist + 1);
        }
    }
}

Python

class Solution(object):
    def distanceK(self, root, target, K):
        ans = []

        # Return distance from node to target if exists, else -1
        # Vertex distance: the # of vertices on the path from node to target
        def dfs(node):
            if not node:
                return -1
            elif node is target:
                subtree_add(node, 0)
                return 1
            else:
                L, R = dfs(node.left), dfs(node.right)
                if L != -1:
                    if L == K: ans.append(node.val)
                    subtree_add(node.right, L + 1)
                    return L + 1
                elif R != -1:
                    if R == K: ans.append(node.val)
                    subtree_add(node.left, R + 1)
                    return R + 1
                else:
                    return -1

        # Add all nodes 'K - dist' from the node to answer.
        def subtree_add(node, dist):
            if not node:
                return
            elif dist == K:
                ans.append(node.val)
            else:
                subtree_add(node.left, dist + 1)
                subtree_add(node.right, dist + 1)

        dfs(root)
        return ans

复杂度分析

时间复杂度： O(N)，其中 N 树的大小。

空间复杂度： O(N)。