Problem

You have a set which contains all positive integers [1, 2, 3, 4, 5, ...].

Implement the SmallestInfiniteSet class:

  • SmallestInfiniteSet() Initializes the SmallestInfiniteSet object to contain all positive integers.
  • int popSmallest() Removes and returns the smallest integer contained in the infinite set.
  • void addBack(int num) Adds a positive integer num back into the infinite set, if it is not already in the infinite set.

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

**Input**
["SmallestInfiniteSet", "addBack", "popSmallest", "popSmallest", "popSmallest", "addBack", "popSmallest", "popSmallest", "popSmallest"]
[[], [2], [], [], [], [1], [], [], []]
**Output**
[null, null, 1, 2, 3, null, 1, 4, 5]

**Explanation**
SmallestInfiniteSet smallestInfiniteSet = new SmallestInfiniteSet();
smallestInfiniteSet.addBack(2);    // 2 is already in the set, so no change is made.
smallestInfiniteSet.popSmallest(); // return 1, since 1 is the smallest number, and remove it from the set.
smallestInfiniteSet.popSmallest(); // return 2, and remove it from the set.
smallestInfiniteSet.popSmallest(); // return 3, and remove it from the set.
smallestInfiniteSet.addBack(1);    // 1 is added back to the set.
smallestInfiniteSet.popSmallest(); // return 1, since 1 was added back to the set and
                                   // is the smallest number, and remove it from the set.
smallestInfiniteSet.popSmallest(); // return 4, and remove it from the set.
smallestInfiniteSet.popSmallest(); // return 5, and remove it from the set.

Constraints

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

Solution

Method 1 – Min-Heap and Set for Efficient Tracking

Intuition

We need to efficiently return and remove the smallest number, and allow numbers to be added back. A min-heap (priority queue) can store numbers that have been added back, and a set can track which numbers are currently in the heap. We also keep a pointer to the next smallest number that hasn’t been popped or added back.

Approach

  1. Use a min-heap to store numbers that have been added back and are smaller than the current pointer.
  2. Use a set to avoid duplicates in the heap.
  3. Keep a pointer curr to the next smallest number not yet popped or added back.
  4. For popSmallest(), if the heap is not empty, pop from the heap; otherwise, return curr and increment it.
  5. For addBack(num), if num is less than curr and not already in the heap, add it to the heap and set.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

#include <queue>
#include <unordered_set>
class SmallestInfiniteSet {
    int curr;
    std::priority_queue<int, std::vector<int>, std::greater<int>> heap;
    std::unordered_set<int> inHeap;
public:
    SmallestInfiniteSet() : curr(1) {}
    int popSmallest() {
        if (!heap.empty()) {
            int x = heap.top(); heap.pop(); inHeap.erase(x); return x;
        }
        return curr++;
    }
    void addBack(int num) {
        if (num < curr && !inHeap.count(num)) {
            heap.push(num); inHeap.insert(num);
        }
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

import java.util.*;
class SmallestInfiniteSet {
    private int curr;
    private PriorityQueue<Integer> heap;
    private Set<Integer> inHeap;
    public SmallestInfiniteSet() {
        curr = 1;
        heap = new PriorityQueue<>();
        inHeap = new HashSet<>();
    }
    public int popSmallest() {
        if (!heap.isEmpty()) {
            int x = heap.poll(); inHeap.remove(x); return x;
        }
        return curr++;
    }
    public void addBack(int num) {
        if (num < curr && !inHeap.contains(num)) {
            heap.offer(num); inHeap.add(num);
        }
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

import heapq
class SmallestInfiniteSet:
    def __init__(self):
        self.curr = 1
        self.heap = []
        self.inHeap = set()
    def popSmallest(self) -> int:
        if self.heap:
            x = heapq.heappop(self.heap)
            self.inHeap.remove(x)
            return x
        res = self.curr
        self.curr += 1
        return res
    def addBack(self, num: int) -> None:
        if num < self.curr and num not in self.inHeap:
            heapq.heappush(self.heap, num)
            self.inHeap.add(num)

Complexity

  • ⏰ Time complexity: O(log n) per operation (heap push/pop).
  • 🧺 Space complexity: O(n) for the heap and set (at most 1000 elements).