题目描述

给定一棵 BST（二叉搜索树），以及其中的一个节点target和一个整数K，求距离target为K的所有节点的值的和。

示例

输入： root = [4,2,6,1,3,5,7], target = 2, K = 1 输出： [1,3,5,7] 解释： 节点 1 和 3 距离 target 的距离为 2，即 K = 2，节点 5 和 7 距离 target 的距离为 3，即 K = 3。 输出距离 target 为 K 的所有节点的值的和为 16。

解法

本题可以采用 DFS（深度优先搜索）或 BFS（广度优先搜索）。

DFS

对于一个根节点root，我们可以计算该节点到目标节点target的最短距离dist，如果dist>K，则返回。否则，分别计算左右子节点的距离，并更新距离值。需要注意的是，当我们更新距离值的时候，需要传递该节点距离目标节点的距离，方便后来查找其子节点。

class Solution(object):
    def __init__(self):
        self.parents = {}

    def distanceK(self, root, target, k):
        self.get_parents(root, None)
        res = []
        self.dfs(target, None, 0, k, res)
        return res

    def get_parents(self, node, parent):
        if not node:
            return
        self.parents[node] = parent
        self.get_parents(node.left, node)
        self.get_parents(node.right, node)

    def dfs(self, node, from_node, dist, k, res):
        if not node:
            return
        if dist == k:
            res.append(node.val)
            return
        if node.left != from_node:
            self.dfs(node.left, node, dist+1, k, res)
        if node.right != from_node:
            self.dfs(node.right, node, dist+1, k, res)
        if self.parents[node] != from_node:
            self.dfs(self.parents[node], node, dist+1, k, res)

BFS

我们可以从目标节点target开始，向四周扩散，直到距离为K。这个过程需要记录节点的父节点，以避免回流。

class Solution(object):
    def distanceK(self, root, target, k):
        q = [(target,0)]
        visited = {target}
        parents = {}

        # 通过 DFS 预处理所有节点的父节点信息
        def dfs(node, parent):
            if node:
                parents[node] = parent
                dfs(node.left, node)
                dfs(node.right, node)
        dfs(root, None)

        # BFS扩散
        while q:
            node, dist = q.pop(0)
            if dist == k:
                return [node.val]
            for next_node in [node.left, node.right, parents[node]]:
                if next_node and next_node not in visited:
                    visited.add(next_node)
                    q.append((next_node, dist+1))
        return []

总结

本题可以采用 DFS 或 BFS 来求解，需要注意的是深度或距离值更新的方法。同时，我们需要预处理出所有节点的父节点信息，以方面后来进行 BFS 扩散。