Problem Overview
LeetCode 2336 presents a SmallestInfiniteSet class that conceptually contains all positive integers. You must support two operations efficiently: popSmallest removes and returns the smallest integer currently in the set, and addBack re-inserts a number that was previously removed.
The naive mental model — an actual set containing all positive integers — is impossible in memory. The problem pushes you toward a smarter representation that captures the infinite natural sequence compactly while still allowing efficient insertions and deletions.
This is a design problem, not a pure algorithm problem. The key skill tested is translating an abstract infinite structure into a finite, efficient data structure — a pattern that shows up in operating system design, database internals, and real-time streaming systems.
- SmallestInfiniteSet initially contains all positive integers {1, 2, 3, 4, ...}
- popSmallest() removes and returns the current smallest integer in the set
- addBack(num) re-adds num to the set — it is a no-op if num is already present
- Constraints: addBack values are in [1, 1000]; at most 1000 calls total
Counter + Sorted Set — The Core Insight
The key observation is that the infinite set always starts as the natural numbers in order. The only numbers that diverge from this sequence are those that have been popped and then added back. Everything from the current frontier upward is still in its natural order.
Maintain a variable currentSmallest starting at 1. This pointer represents the smallest natural number not yet popped that has never been re-added. Alongside it, keep a sorted set (SortedSet in Java, SortedList in Python) that tracks numbers which were popped earlier and then added back via addBack — specifically, numbers that are strictly less than currentSmallest.
When popSmallest is called, check the sorted set first. If it is non-empty, its minimum is smaller than currentSmallest, so return and remove that minimum. If the sorted set is empty, the next smallest is currentSmallest itself — return it and increment the counter.
Design Insight
The counter handles infinite natural progression without storing anything. The sorted set only holds "exceptions" — re-added numbers below the counter. This keeps memory bounded to the number of addBack calls, not the size of the conceptual infinite set.
Pop Operation — Sorted Set First, Counter Second
The popSmallest operation is a two-branch decision tree. First, inspect the sorted set of re-added numbers. If it contains any elements, the minimum of that set is guaranteed to be smaller than currentSmallest (by the invariant we maintain in addBack). Remove that minimum from the set and return it.
If the sorted set is empty, no previously-popped number has been re-added below the frontier. The smallest element is therefore currentSmallest. Return its value and increment currentSmallest by one, advancing the frontier.
This two-branch design guarantees correctness: the set always holds exceptions below the counter, and the counter always represents the start of the unmodified natural sequence. Combining them gives the global minimum at any point in time.
- 1Check if addedBack sorted set is non-empty
- 2If non-empty: remove and return the minimum element from addedBack
- 3If empty: capture currentSmallest, increment currentSmallest, return captured value
- 4Both paths run in O(log n) where n is the size of the addedBack set
AddBack Operation — Only Re-add What Was Removed
The addBack(num) operation must respect the infinite set semantics: if num is already in the set, the call is a no-op. Numbers are in the set if they are >= currentSmallest (they were never popped) or if they are already present in the addedBack sorted set.
Therefore: if num >= currentSmallest, num is still in the natural frontier — do nothing. If num < currentSmallest, the number was previously popped at some point. Add it to the addedBack sorted set only if it is not already there. The sorted set data structure naturally handles deduplication.
This conditional check is crucial. Without it, adding a number >= currentSmallest to the set would corrupt the invariant — the counter already represents all numbers from currentSmallest upward, so inserting one of them into addedBack would cause it to be returned twice.
Deduplication Guarantee
The sorted set enforces uniqueness automatically: addBack(5) called twice inserts 5 only once. When popSmallest later pulls 5 from the set, it appears exactly once — no double-counting. This is why using a set instead of a list or heap alone is critical.
Walk-Through Example — Tracing the State
Walking through a concrete sequence of operations builds the intuition for why the counter-plus-set invariant is always maintained. Each step below shows the state of currentSmallest and addedBack after the operation.
Notice how the sorted set only ever holds numbers smaller than currentSmallest. When addBack(1) is called after the counter has advanced past 1, it correctly inserts 1 into the set. When addBack(2) is called later, it inserts 2 into the set — and the next pop draws from the set rather than the counter.
The invariant "addedBack contains only numbers < currentSmallest" is maintained after every operation. This is the key to correctness: the set and counter partition responsibility cleanly, with no overlap.
- 1pop() → return 1, currentSmallest becomes 2, addedBack = {}
- 2pop() → return 2, currentSmallest becomes 3, addedBack = {}
- 3addBack(1) → 1 < 3, insert into set, addedBack = {1}
- 4pop() → addedBack non-empty, remove and return 1, addedBack = {}
- 5pop() → addedBack empty, return 3, currentSmallest becomes 4
- 6addBack(2) → 2 < 4, insert into set, addedBack = {2}
- 7pop() → addedBack non-empty, remove and return 2, addedBack = {}
Code Walkthrough — Python and Java
In Python, use sortedcontainers.SortedList for the addedBack collection. Initialize currentSmallest = 1 and addedBack = SortedList(). In popSmallest: if addedBack, return addedBack.pop(0); else increment and return currentSmallest. In addBack: if num < self.currentSmallest and num not in addedBack, addedBack.add(num). Both operations run in O(log n).
In Java, use a TreeSet<Integer> for addedBack. TreeSet provides O(log n) add, remove, and first (minimum) operations with automatic deduplication. In popSmallest: if !addedBack.isEmpty(), return addedBack.pollFirst(); else return currentSmallest++. In addBack: if num < currentSmallest, addedBack.add(num) — TreeSet ignores duplicates.
The space complexity is O(k) where k is the number of distinct numbers in addedBack at any time — bounded by the number of addBack calls. Time complexity per operation is O(log k). Given the 1000-call constraint, the worst-case set size is 1000 elements, making log k approximately 10.
Python Pitfall
Python doesn't have a built-in SortedSet. Use sortedcontainers.SortedList or maintain a heap (heapq) combined with a set for deduplication. A plain set gives O(1) add but no O(1) minimum — you'd need O(n) min() each time, which defeats the purpose.