Smallest Number in 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 SmallestInfiniteSet:
    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 SmallestInfiniteSet {
  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 SmallestInfiniteSet {
    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 SmallestInfiniteSet {
    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);
        }
    }
};

Explanation

The set is infinite, so we cannot actually store every positive integer. The clever idea is to split the set into two parts: a single counter next for the long unbroken run of numbers that were never popped, and a min-heap for the smaller numbers that were popped and then added back.

The counter next remembers the smallest number that has never been popped (it starts at 1). Everything from next upward is still in the set without needing to be stored individually.

For popSmallest(), we first check the heap. If it holds any added-back numbers, the smallest of those is even smaller than next, so we pop and return it. Otherwise we return next and bump the counter up by one.

For addBack(x), we only do something if x was actually popped before (meaning x < next) and is not already waiting in the heap. We use a set to avoid duplicate pushes.

Example: pop returns 1, 2, 3, so next is now 4. Then addBack(2) drops 2 into the heap. The next pop sees 2 in the heap (smaller than 4) and returns 2.

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