< Back





2336. Smallest Number in Infinite Set

We're asked to create a class that will simulate that we have an infinite set of numbers, starting at
1. The class should support the operation to pop the smallest number from the infinite set. Once a
number is popped from the infinite set, it's considered to be removed from the set. The class should
also support adding a number back into the infinite set, iff it was previously removed.

To make this happen, we use a min heap, a set, and a counter to keep track of the numbers we can add
to the heap. If the heap exists, it means numbers were added back to the infinite set, so we pop the
smallest number from the heap. Otherwise, we return the current number and increment the counter. We
also keep track of what numbers are current removed from the set.

When a number is added back to the set, we check if it's in the removed set. If it is, we remove it
from the set and add it to the heap.

The solution is as follows:


  from heapq import heappop, heappush

  class SmallestInfiniteSet:
      def __init__(self):
          self.h = []
          self.curr = 1
          self.out = set()

      def popSmallest(self) -> int:
          if self.h:
              res = heappop(self.h)
          else:
              res, self.curr = self.curr, self.curr + 1
          
          self.out.add(res)
          return res

      def addBack(self, num: int) -> None:
          if num in self.out:
              self.out.discard(num)
              heappush(self.h, num)


_ Time Complexity:

  O((m + n)*log(n)) - For m popSmallest operations and n addBack operations, the heap operations take
  O(log(n)) time.

_ Space Complexity:

  O(n) - The heap can store up to n elements in the worst case.