Problem statement

A set starts out holding every positive integer: 1, 2, 3, 4, …. Write the class SmallestInfiniteSet:

• SmallestInfiniteSet() creates the set with all positive integers in it.
• int popSmallest() removes the smallest number in the set and returns it.
• void addBack(int num) puts num back into the set if it is not already there (otherwise nothing happens).

The judge calls the methods in the order listed in the first input line, with the arguments in the second; the expected output holds each call's result, null for the constructor and for every method that returns nothing.

Example 1

Input: ["SmallestInfiniteSet","addBack","popSmallest","popSmallest","popSmallest","addBack","popSmallest","popSmallest","popSmallest"]
[[],[2],[],[],[],[1],[],[],[]]
Output: [null,null,1,2,3,null,1,4,5]
Explanation: 2 is already in the set, so addBack(2) changes nothing. The pops take 1, 2 and 3; 1 comes back and is the next smallest, followed by 4 and 5.

Example 2

Input: ["SmallestInfiniteSet","popSmallest","popSmallest","popSmallest","addBack","addBack","addBack","popSmallest","popSmallest","popSmallest"]
[[],[],[],[],[3],[1],[3],[],[],[]]
Output: [null,1,2,3,null,null,null,1,3,4]
Explanation: After 1, 2 and 3 are taken, 3 and 1 are put back (the second addBack(3) does nothing). The pops return 1, 3, then 4 — 2 is still gone.

Constraints

• 1 <= num <= 1000
• At most 1000 calls are made in total to popSmallest and addBack

Hints for Smallest Number in Infinite Set

1. Hint 1

   Almost all of the set is one unbroken run: every number from some next upwards is still there. Track next instead of the infinite set.

2. Hint 2

   The only other members are numbers smaller than next that were added back. Keep those in a structure that gives you its minimum quickly.

3. Hint 3

   popSmallest takes the smallest added-back number if there is one (it is below next), otherwise returns next and advances it. addBack(num) matters only when num < next and num is not already stored.

How to solve Smallest Number in Infinite Set

Split the set into two parts: the infinite tail next, next+1, … (never touched yet) and a small sorted collection of numbers below next that were added back. The minimum is always the smallest added-back number if any exist, else next.

Approach

1. Start with next = 1 and an empty ordered set (or min-heap plus membership set) back.
2. popSmallest(): if back is non-empty, remove and return its minimum; otherwise return next and increase it by one.
3. addBack(num): if num < next and num is not in back, insert it. A number >= next is already in the set.

Why it works

The invariant is that the set equals back ∪ {next, next+1, …} with every element of back below next. A pop of the tail keeps it (the tail just starts one later); a pop from back keeps it; addBack inserts only numbers missing from both parts. Since all of back lies below next, its minimum, when it exists, is the set's minimum.

Complexity

• Time — O(log n) per call with a heap or balanced tree (O(n) insertion with a sorted array), n ≤ 1000
• Space — O(n)

Pitfalls

• Adding back a number that is still present must not create a duplicate — check num < next and membership in back.
• A plain min-heap without a membership check returns the same number twice after a double addBack.

Reference solution

Python

import heapq

class SmallestInfiniteSet:

    def __init__(self):
        self.next = 1      # every number >= next is in the set
        self.heap = []     # added-back numbers below next
        self.inHeap = set()

    def popSmallest(self) -> int:
        if self.heap:
            x = heapq.heappop(self.heap)
            self.inHeap.discard(x)
            return x
        self.next += 1
        return self.next - 1

    def addBack(self, num: int) -> None:
        if num < self.next and num not in self.inHeap:
            self.inHeap.add(num)
            heapq.heappush(self.heap, num)

Java

class SmallestInfiniteSet {
    private int next = 1; // every number >= next is in the set
    private TreeSet<Integer> back = new TreeSet<>(); // added-back numbers below next

    public SmallestInfiniteSet() {
    }

    public int popSmallest() {
        if (!back.isEmpty()) return back.pollFirst();
        return next++;
    }

    public void addBack(int num) {
        if (num < next) back.add(num);
    }
}

C++

class SmallestInfiniteSet {
    int next = 1;  // every number >= next is in the set
    set<int> back; // added-back numbers below next
public:
    SmallestInfiniteSet() {}

    int popSmallest() {
        if (!back.empty()) {
            int x = *back.begin();
            back.erase(back.begin());
            return x;
        }
        return next++;
    }

    void addBack(int num) {
        if (num < next) back.insert(num);
    }
};

JavaScript

var SmallestInfiniteSet = function() {
    this.next = 1;  // every number >= next is in the set
    this.back = []; // added-back numbers below next, sorted ascending
};

SmallestInfiniteSet.prototype.popSmallest = function() {
    if (this.back.length > 0) return this.back.shift();
    return this.next++;
};

SmallestInfiniteSet.prototype.addBack = function(num) {
    if (num >= this.next) return;
    var lo = 0, hi = this.back.length;
    while (lo < hi) {
        var mid = (lo + hi) >> 1;
        if (this.back[mid] < num) lo = mid + 1; else hi = mid;
    }
    if (lo < this.back.length && this.back[lo] === num) return;
    this.back.splice(lo, 0, num);
};

Also on the editorial tab: C, C#, Go, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

More hash table problems

• Implement Router Medium
• Design Task Manager Medium
• Frequency Tracker Medium
• Find Consecutive Integers from a Data Stream Medium
• Design Memory Allocator Medium
• Longest Uploaded Prefix Medium