2336. 无限集中的最小数字 : 容易又高效的分类做法

题目描述

这是 LeetCode 上的 「2336. 无限集中的最小数字」 ，难度为 「中等」。

Tag : 「优先队列（堆）」、「哈希表」

现有一个包含所有正整数的集合。

实现 SmallestInfiniteSet 类：

SmallestInfiniteSet() 初始化 SmallestInfiniteSet 对象以包含所有正整数。
int popSmallest() 移除并返回该无限集中的最小整数。
void addBack(int num) 如果正整数 num 不存在于无限集中，则将一个 num 添加到该无限集中。

示例：

输入
["SmallestInfiniteSet", "addBack", "popSmallest", "popSmallest", "popSmallest", "addBack", "popSmallest", "popSmallest", "popSmallest"]
[[], [2], [], [], [], [1], [], [], []]

输出
[null, null, 1, 2, 3, null, 1, 4, 5]

解释
SmallestInfiniteSet smallestInfiniteSet = new SmallestInfiniteSet();
smallestInfiniteSet.addBack(2);    // 2 已经在集合中，所以不做任何变更。
smallestInfiniteSet.popSmallest(); // 返回 1 ，因为 1 是最小的整数，并将其从集合中移除。
smallestInfiniteSet.popSmallest(); // 返回 2 ，并将其从集合中移除。
smallestInfiniteSet.popSmallest(); // 返回 3 ，并将其从集合中移除。
smallestInfiniteSet.addBack(1);    // 将 1 添加到该集合中。
smallestInfiniteSet.popSmallest(); // 返回 1 ，因为 1 在上一步中被添加到集合中，
                                   // 且 1 是最小的整数，并将其从集合中移除。
smallestInfiniteSet.popSmallest(); // 返回 4 ，并将其从集合中移除。
smallestInfiniteSet.popSmallest(); // 返回 5 ，并将其从集合中移除。

提示：

最多调用 popSmallest 和 addBack 方法共计次

优先队列（小根堆）+ 哈希表

使用 idx 代表顺序弹出的集合左边界，范围内的数均为待弹出，起始有。

考虑当调用 addBack 往集合添加数值 x 时，该如何处理：

：数值本身就存在于集合中，忽略该添加操作；
：数值刚好位于边界左侧，更新；
：考虑将数值添加到某个容器中，该容器支持返回最小值，容易联想到“小根堆”；但小根堆并没有“去重”功能，为防止重复弹出，还需额外使用“哈希表”来记录哪些元素在堆中。

该做法本质上将集合分成两类：一类是从 idx 到正无穷的连续段，对此类操作的复杂度为；一类是比 idx 要小的离散类数集，对该类操作复杂度为，其中为调用 addBack 的最大次数。

Java 代码：

class SmallestInfiniteSet {
    boolean[] vis = new boolean[1010];
    PriorityQueue<Integer> q = new PriorityQueue<>((a,b)->a-b);
    int idx = 1;
    public int popSmallest() {
        int ans = -1;
        if (!q.isEmpty()) {
            ans = q.poll();
            vis[ans] = false;
        } else {
            ans = idx++;
        }
        return ans;
    }
    public void addBack(int x) {
        if (x >= idx || vis[x]) return ;
        if (x == idx - 1) {
            idx--;
        } else {
            q.add(x);
            vis[x] = true;
        }
    }
}

C++ 代码：

class SmallestInfiniteSet {
public:
    vector<bool> vis;
    priority_queue<int, vector<int>, greater<int>> q;
    int idx;
    SmallestInfiniteSet() : idx(1) {
        vis.resize(1010, false);
    }
    int popSmallest() {
        int ans = -1;
        if (!q.empty()) {
            ans = q.top();
            q.pop();
            vis[ans] = false;
        } else {
            ans = idx++;
        }
        return ans;
    }
    void addBack(int x) {
        if (x >= idx || vis[x]) return;
        if (x == idx - 1) {
            idx--;
        } else {
            q.push(x);
            vis[x] = true;
        }
    }
};

Python 代码：

class SmallestInfiniteSet:
    def __init__(self):
        self.vis = [False] * 1010
        self.q = []
        self.idx = 1

    def popSmallest(self):
        ans = -1
        if self.q:
            ans = heapq.heappop(self.q)
            self.vis[ans] = False
        else:
            ans = self.idx
            self.idx += 1
        return ans

    def addBack(self, x):
        if x >= self.idx or self.vis[x]:
            return
        if x == self.idx - 1:
            self.idx -= 1
        else:
            heapq.heappush(self.q, x)
            self.vis[x] = True