寻找二叉搜索树中的目标节点

标签: 深度优先搜索 二叉搜索树 二叉树

难度: Easy

某公司组织架构以二叉搜索树形式记录,节点值为处于该职位的员工编号。请返回第 cnt 大的员工编号。

示例 1:

输入:root = [7, 3, 9, 1, 5], cnt = 2
       7
      / \
     3   9
    / \
   1   5
输出:7

示例 2:

输入: root = [10, 5, 15, 2, 7, null, 20, 1, null, 6, 8], cnt = 4
       10
      / \
     5   15
    / \    \
   2   7    20
  /   / \ 
 1   6   8
输出: 8

提示:

  • 1 ≤ cnt ≤ 二叉搜索树元素个数

Submission

运行时间: 60 ms

内存: 18.7 MB

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

class Solution:

    k = None
    res = None

    def kthLargest(self, root: TreeNode, k: int) -> int:
        self.k = k

        def dfs(root):
            if root is None:
                return
            dfs(root.right)
            if self.k == 0:
                return
            self.k -= 1
            if self.k == 0:
                self.res = root.val
            dfs(root.left)

        dfs(root)
        return self.res

Explain

题解利用了二叉搜索树的性质,即其中序遍历结果是递增序列。要找第 cnt 大的元素,可以进行一种反向中序遍历(即先访问右子树,再访问根,最后访问左子树)。这样第 cnt 次访问的节点就是第 cnt 大的节点。解题关键在于通过计数器 self.k 来记录访问的元素数,当计数器减到0时,当前节点即为所求的第 cnt 大的元素。

时间复杂度: O(n)

空间复杂度: O(n)

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

class Solution:

    k = None
    res = None

    def kthLargest(self, root: TreeNode, k: int) -> int:
        self.k = k  # 初始化计数器

        def dfs(root):
            if root is None:  # 如果节点为空,返回
                return
            dfs(root.right)  # 首先访问右子节点
            if self.k == 0:  # 如果计数器为0,提前结束
                return
            self.k -= 1  # 访问一个节点,计数器减1
            if self.k == 0:  # 如果计数器减到0,当前节点即为所求
                self.res = root.val
            dfs(root.left)  # 最后访问左子节点

        dfs(root)  # 开始递归
        return self.res  # 返回结果

Explore

在反向中序遍历中,通过每访问一个节点就将`self.k`的值减1的方式来更新。这种方法确保了每次访问节点时,计数器都会反映剩余需要访问的节点数量。当`self.k`减到0时,表示已经找到了第`k`大的元素,因此这种更新方式能够精确控制遍历的深度和范围,达到快速定位目标节点的目的。

在`self.k`为0时立即停止遍历是为了提高算法的效率。因为一旦找到第`k`大的元素,继续遍历其他节点将不再有实际意义,只会增加不必要的计算和时间消耗。因此,提前终止遍历可以节省资源,使算法更加高效。

当树为空(根节点为null),`dfs`函数将在一开始就返回,因为不会有任何节点可访问。此时,`self.res`将保持为None或其初始值,因此最终返回的结果将是None。同样,如果输入的`k`值大于树中的节点总数,遍历完成后`self.k`仍大于0,这意味着没有足够的节点来满足第`k`大的要求,最终输出结果也将是None。

是的,在实际应用中应该添加对`k`值非法情况的检测。这包括确保`k`不为负数、不为零,并且不超过树中节点的总数。这样的错误检查可以防止运行时错误和不合逻辑的输出,提高程序的健壯性和用户体验。