程序员面试金典 - 面试题 10.10. 数字流的秩

题目难度: 中等

原题链接

今天继续更新程序员面试金典系列, 大家在公众号 算法精选 里回复 面试金典 就能看到该系列当前连载的所有文章了, 记得关注哦~

题目描述

假设你正在读取一串整数。每隔一段时间，你希望能找出数字 x 的秩(小于或等于 x 的值的个数)。请实现数据结构和算法来支持这些操作，也就是说：

• 实现 track(int x) 方法，每读入一个数字都会调用该方法；

• 实现 getRankOfNumber(int x) 方法，返回小于或等于 x 的值的个数。

• 注意：本题相对原题稍作改动

示例:

• 输入:
  • [“StreamRank”, “getRankOfNumber”, “track”, “getRankOfNumber”]
  • [[], [1], [0], [0]]
• 输出:
  • [null,0,null,1]

提示：

• x <= 50000
• track 和 getRankOfNumber 方法的调用次数均不超过 2000 次

题目思考

1. 如何尽可能提高 track 和 getRankOfNumber 方法的效率?

解决方案

思路

• 分析题目, 要实现插入和查询操作, 一个比较简单的思路是: 维护一个有序数组, 插入时利用二分查找找到合适的插入位置, 继续保持有序性; 而查询时利用二分查找得到小于等于 x 的元素数目
• 这个思路虽然查询的复杂度较低 (O(logN)), 但插入元素时的复杂度较高 (O(N)), 因为每次插入时都需要将插入位置之后的元素往后移动
• 那有哪些数据结构可以同时兼顾插入和查询呢?
• 一个比较合适的数据结构就是二叉搜索树, 它可以做到插入和查询平均都是 O(logN) 的复杂度
• 但这道题目不一样的地方在于, 查询时查找的是所有小于等于 x 的数目, 而常规的二叉搜索树不能做到这一点, 所以我们需要进行一些改进
• 重新分析二叉搜索树, 我们不难发现它有一个很重要的性质, 就是左子树的所有节点都小于当前节点
• 我们可以利用这一点, 对每个节点额外维护一个 lcnt, 代表它及其左子树的数字数目之和
• 然后查询的时候, 对于每个遍历到的节点, 如果其节点值小于等于 x, 就累加这一 lcnt, 最终遍历结束时累加的结果自然就是所有小于等于 x 的数目了
• 而插入的时候, 如果需要继续往当前节点的左边查找插入位置, 也需要将当前节点的 lcnt 加 1
• 下面的代码就对应了上面的整个过程, 并且有详细的注释, 方便大家理解

复杂度

• 时间复杂度平均 O(logN) 最坏 O(N): 插入和查找二叉搜索树所需遍历的节点数目是层高, 平均层高为 logN, 而最差情况下的层高是 N (插入的数字是顺序或逆序的, 此时一直向同一侧增加层数)
• 空间复杂度 O(N): 需要存储所有数字

代码

class StreamRank:
    # 二叉搜索树+额外存储左子树数目
    class Node:
        def __init__(self, val) -> None:
            self.val = val
            self.left = None
            self.right = None
            # lcnt存储当前节点自身及左子树的数目之和
            self.lcnt = 1

    def __init__(self):
        # 初始化根节点为最大值, 作为哨兵节点
        self.root = self.Node(float("inf"))

    def track(self, x: int) -> None:
        # 查找x应该插入的位置, 并在需要时更新父节点的lcnt
        # pre初始化为哨兵节点root, cur则初始化为其左子节点 (因为root是最大值, 所有节点都在其左子树上)
        pre, cur = self.root, self.root.left
        while cur:
            if cur.val == x:
                # 当前节点值等于x, 无需增加新节点, 只需将当前节点的lcnt加1即可
                cur.lcnt += 1
                # 然后直接退出, 因为其左右子树的lcnt都无需变化, 且无需插入新节点
                return
            elif cur.val > x:
                # 当前节点值大于x, 将其lcnt加1, 并向左子树移动
                cur.lcnt += 1
                pre, cur = cur, cur.left
            else:
                # 当前节点值小于x, 需要向右子树移动
                # 新节点不在左子树上, 所以lcnt也不能增加
                pre, cur = cur, cur.right

        # 此时根据父节点pre和x的关系来插入新节点
        if pre.val > x:
            # 父节点大于x, 插入x作为父节点左儿子
            pre.left = self.Node(x)
        else:
            # 父节点小于x, 插入x作为父节点右儿子
            pre.right = self.Node(x)

    def getRankOfNumber(self, x: int) -> int:
        # 利用二叉搜索树的有序性查找x并累加小于等于x的数目
        cur = self.root
        res = 0
        while cur:
            if cur.val > x:
                # 当前节点值更大, 向左子树查找, 同时不能累加当前节点的lcnt
                cur = cur.left
            else:
                # 当前节点值小于等于x, 则其左子树部分都小于等于x, 也即累加当前节点的lcnt
                res += cur.lcnt
                if cur.val == x:
                    # 当前节点的值恰好等于x, 无需继续查找, 直接跳出循环
                    break
                else:
                    # 当前节点的值小于x, 继续向其右子树查找
                    cur = cur.right
        # 最终结果res即为树中小于等于x的数目
        return res

---

大家可以在下面这些地方找到我~

我的 GitHub

我的 Leetcode

我的 CSDN

我的知乎专栏

我的头条号

我的牛客网博客

我的公众号: 算法精选, 欢迎大家扫码关注~