二叉搜索树最近节点查询

标签: 树 深度优先搜索 二叉搜索树 数组 二分查找 二叉树

难度: Medium

给你一个 二叉搜索树 的根节点 root ，和一个由正整数组成、长度为 n 的数组 queries 。

请你找出一个长度为 n 的 二维 答案数组 answer ，其中 answer[i] = [mini, maxi] ：

  • mini 是树中小于等于 queries[i] 的 最大值 。如果不存在这样的值，则使用 -1 代替。
  • maxi 是树中大于等于 queries[i] 的 最小值 。如果不存在这样的值，则使用 -1 代替。

返回数组 answer 。

示例 1 ：

输入：root = [6,2,13,1,4,9,15,null,null,null,null,null,null,14], queries = [2,5,16]
输出：[[2,2],[4,6],[15,-1]]
解释：按下面的描述找出并返回查询的答案：
- 树中小于等于 2 的最大值是 2 ，且大于等于 2 的最小值也是 2 。所以第一个查询的答案是 [2,2] 。
- 树中小于等于 5 的最大值是 4 ，且大于等于 5 的最小值是 6 。所以第二个查询的答案是 [4,6] 。
- 树中小于等于 16 的最大值是 15 ，且大于等于 16 的最小值不存在。所以第三个查询的答案是 [15,-1] 。

示例 2 ：

输入：root = [4,null,9], queries = [3]
输出：[[-1,4]]
解释：树中不存在小于等于 3 的最大值，且大于等于 3 的最小值是 4 。所以查询的答案是 [-1,4] 。

提示：

  • 树中节点的数目在范围 [2, 105] 内
  • 1 <= Node.val <= 106
  • n == queries.length
  • 1 <= n <= 105
  • 1 <= queries[i] <= 106

Submission

运行时间: 1048 ms

内存: 153.1 MB

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def closestNodes(self, root: Optional[TreeNode], queries: List[int]) -> List[List[int]]:
        vals = []
        # 通过中序遍历BST构建有序数组
        def dfs(node):
            if not node: return
            dfs(node.left)
            vals.append(node.val)
            dfs(node.right)
        dfs(root)

        res = []
        n = len(vals)
        for q in queries:
            j = bisect_left(vals, q) # >= q的最小索引
            max_q = vals[j] if j < n else -1
            i = bisect_left(vals, q + 1) - 1 # <= q的最大索引
            min_q = vals[i] if i >= 0 else -1
            res.append([min_q, max_q])
        return res

Explain

该题解首先通过对给定的二叉搜索树进行中序遍历，得到一个有序数组，该数组按从小到大的顺序包含了树中所有的元素。中序遍历对于二叉搜索树的特性是得到的遍历结果是有序的。接着，对于每一个查询值，使用二分查找确定该查询值在有序数组中的位置，以及最接近该查询值的元素。具体来说，使用`bisect_left`函数查找大于等于查询值的最小元素，以及查找小于等于查询值的最大元素。这样可以分别确定每个查询的`maxi`和`mini`值。

时间复杂度: O(m + n log m)

空间复杂度: O(m + n)

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def closestNodes(self, root: Optional[TreeNode], queries: List[int]) -> List[List[int]]:
        vals = []
        # 通过中序遍历BST构建有序数组
        def dfs(node):
            if not node: return
            dfs(node.left)  # 遍历左子树
            vals.append(node.val)  # 访问当前节点
            dfs(node.right)  # 遍历右子树
        dfs(root)

        res = []
        n = len(vals)
        for q in queries:
            j = bisect_left(vals, q)  # 找到大于等于q的最小值的位置
            max_q = vals[j] if j < n else -1  # 如果位置有效，则为maxi，否则为-1
            i = bisect_left(vals, q + 1) - 1  # 找到小于等于q的最大值的位置
            min_q = vals[i] if i >= 0 else -1  # 如果位置有效，则为mini，否则为-1
            res.append([min_q, max_q])
        return res

Explore

在处理查询时，我们需要找到数组中小于等于查询值`q`的最大值和大于等于查询值`q`的最小值。使用`bisect_left`查找`q`可以直接得到大于等于`q`的最小值的索引，这是因为`bisect_left`返回的是插入点，即第一个不小于`q`的元素的位置。为了找到小于等于`q`的最大值，我们可以查找`q + 1`的插入点然后减一，这样得到的是最后一个不大于`q`的元素的位置，即小于等于`q`的最大值的位置。这种方法利用了`bisect_left`函数的特性来简化查找过程。

中序遍历是一种深度优先搜索方法，它可以正确访问二叉搜索树的每个节点，并且按照左子树-根节点-右子树的顺序进行。即使在极端不平衡的二叉搜索树（如所有节点都倾向于一边的树）中，中序遍历仍然可以保持正确的访问顺序，因为它递归地访问每个子树。确保中序遍历正确执行的关键是正确实现递归逻辑，处理好所有节点的左右子树递归调用，以及在递归过程中正确地访问和记录节点值。

在使用`bisect_left`找到大于等于`q`的最小值的位置后，需要检查得到的索引`j`是否等于数组长度`n`。如果`j`等于`n`，这意味着所有数组中的元素都小于`q`，因此不存在一个大于或等于`q`的元素。在这种情况下，直接将对应的结果设置为`-1`，而不是尝试访问`vals[j]`，这样可以避免数组越界错误。通过这种边界检查，确保程序的健壮性和正确性。