Python二叉搜索树指定距离节点查找：修正仅返回子树节点的代码

问题分析 #

当前代码仅能遍历目标节点的子树方向节点，无法处理向上（祖先节点方向）的距离计算。例如节点3的情况，其距离为2的节点包括父节点的父节点5，以及父节点的另一子节点1，这些都需要额外处理。

修正后完整代码 #

class BST2(BinarySearchTree):
    def find_dist_k(self, n: int, k: int) -> list:
        result = []
        node = self.search(n)
        if node is None:
            return result
        if k < 0:
            print("参数k必须为正整数")
            return result
        
        # 第一步：收集目标节点向下距离k的所有节点
        self._find_down(node, k, result)
        
        # 第二步：收集祖先方向符合距离要求的节点
        path = self._get_root_to_node_path(node)
        # 遍历路径中的每个祖先节点（排除目标节点自身）
        for idx in range(len(path) - 1):
            ancestor = path[idx]
            # 计算当前祖先到目标节点的距离
            dist_to_target = len(path) - 1 - idx
            remaining_k = k - dist_to_target
            
            if remaining_k < 0:
                continue
            
            # 找到祖先的另一个子节点（不是路径中指向目标节点的子节点）
            next_node_in_path = path[idx + 1]
            other_child = ancestor.right if ancestor.left == next_node_in_path else ancestor.left
            
            if remaining_k == 0:
                # 祖先节点本身就是距离为k的节点
                result.append(ancestor.elem)
            else:
                # 遍历另一个子树，找距离为remaining_k-1的节点
                self._find_down(other_child, remaining_k - 1, result)
        
        return result
    
    def _find_down(self, node, k, result):
        """递归遍历节点的子树，收集距离当前节点k的所有节点"""
        if node is None:
            return
        if k == 0:
            result.append(node.elem)
            return
        self._find_down(node.left, k - 1, result)
        self._find_down(node.right, k - 1, result)
    
    def _get_root_to_node_path(self, node):
        """获取从根节点到目标节点的路径列表（顺序为根到目标节点）"""
        path = []
        current = self.root
        while current is not None:
            path.append(current)
            if current.elem == node.elem:
                break
            elif node.elem < current.elem:
                current = current.left
            else:
                current = current.right
        return path

关键修正点说明 #

1. 拆分向下遍历逻辑：将原递归方法重命名为_find_down，专注处理目标节点子树方向的距离查找，保持单一职责。
2. 获取根到目标节点的路径：通过_get_root_to_node_path方法，利用二叉搜索树左小右大的特性遍历得到完整路径，全程无需依赖parent属性。
3. 处理祖先方向节点：
   • 遍历路径中的每个祖先节点，计算该祖先到目标节点的实际距离dist_to_target。
   • 根据剩余距离remaining_k = k - dist_to_target做不同处理：
     • 若remaining_k == 0，直接将该祖先节点加入结果（它就是距离目标节点k的节点）。
     • 若remaining_k > 0，遍历该祖先的另一子树（非目标节点所在的子树），查找距离为remaining_k - 1的节点（因为从祖先到该子节点已经消耗了一步距离）。

测试验证 #

针对输入列表input_list_01 = [5, 12, 2, 1, 3, 9]：

• 调用find_dist_k(3, 2)时，结果为[5, 1]，符合距离定义（3→2→5 距离2；3→2→1 距离2）。
• 调用find_dist_k(5, 2)时，结果仍为[1, 3, 9]，与原预期一致。

内容的提问来源于stack exchange，提问作者user20442513