Smallest Available Number From an Infinite Set

medium heap design

Problem

Implement a data structure containing every positive integer 1, 2, 3, …. Support two operations: popSmallest() removes and returns the smallest still-present integer, and addBack(x) reinserts x if it had previously been popped. Track the smallest never-popped number with a counter, and the previously-popped numbers in a min-heap.

Inputops = pop, pop, pop, addBack 2, pop

Output1, 2, 3, _, 2

After popping 1,2,3 the counter is at 4. Adding back 2 puts 2 in the heap; next pop returns 2.

operations (e.g. "pop; pop; pop; add 2; pop")

import heapq
class SmallestSet:
    def __init__(self):
        self.next = 1
        self.heap = []
        self.in_set = set()
    def pop_smallest(self):
        if self.heap:
            v = heapq.heappop(self.heap); self.in_set.discard(v); return v
        v = self.next; self.next += 1; return v
    def add_back(self, x):
        if x < self.next and x not in self.in_set:
            self.in_set.add(x); heapq.heappush(self.heap, x)

class SmallestSet {
  constructor() { this.next = 1; this.heap = []; this.in = new Set(); }
  popSmallest() {
    if (this.heap.length) { const v = this.heap.shift(); this.in.delete(v); return v; }
    return this.next++;
  }
  addBack(x) {
    if (x < this.next && !this.in.has(x)) {
      this.in.add(x); this.heap.push(x); this.heap.sort((a, b) => a - b);
    }
  }
}

class SmallestSet {
    private int next = 1;
    private PriorityQueue<Integer> heap = new PriorityQueue<>();
    private Set<Integer> in = new HashSet<>();
    public int popSmallest() {
        if (!heap.isEmpty()) { int v = heap.poll(); in.remove(v); return v; }
        return next++;
    }
    public void addBack(int x) {
        if (x < next && !in.contains(x)) {
            in.add(x); heap.offer(x);
        }
    }
}

class SmallestSet {
    int nxt = 1;
    priority_queue<int, vector<int>, greater<int>> heap;
    unordered_set<int> in_set;
public:
    int popSmallest() {
        if (!heap.empty()) { int v = heap.top(); heap.pop(); in_set.erase(v); return v; }
        return nxt++;
    }
    void addBack(int x) {
        if (x < nxt && !in_set.count(x)) {
            in_set.insert(x); heap.push(x);
        }
    }
};

Time: O(log n) per op Space: O(popped not yet returned)