Smallest Number in Infinite Set

Smallest Number in Infinite Set

leetcode link

In this article, we will explore the design and implementation of the SmallestInfiniteSet class. This class allows you to manage an infinite set of positive integers with two primary operations: popping the smallest element and adding elements back into the set. By using a combination of a min-heap and a set, we ensure that these operations are efficient. Let’s begin by looking at the code.

Full Implementation

from heapq import heappush, heappop

class SmallestInfiniteSet:

    def __init__(self):
        self.current = 1  # Tracks the next smallest element
        self.added = []  # Min-heap for managing added-back elements
        self.addedSet = set()  # Set for O(1) checks of added elements

    def popSmallest(self) -> int:
        if not self.added:
            result = self.current
            self.current += 1
        else:
            result = heappop(self.added)
            self.addedSet.remove(result)
        return result

    def addBack(self, num: int) -> None:
        if num >= self.current or num in self.addedSet:
            return
        heappush(self.added, num)
        self.addedSet.add(num)

Why Does This Approach Work?

The design choices in this implementation might not be immediately obvious. The key to understanding the solution lies in the management of two structures: self.current, which represents the next smallest element that hasn’t been popped, and self.added, a min-heap that holds previously popped elements that were added back.

Tracking the Smallest Element with self.current

One crucial design decision is that self.current only increases and is never decremented. This is essential because it means we don’t need to explicitly store every integer from the infinite set. Instead, we increment self.current whenever the smallest element is popped, and it acts as a pointer to the next smallest number.

Example:

  • Initially, self.current = 1.
  • Call popSmallest() returns 1 and updates self.current to 2.
  • The next call to popSmallest() returns 2, and now self.current becomes 3.

This approach ensures we simulate an infinite set without needing to store all the integers, which would be computationally expensive.

Why Is self.current So Important?

If the addBack() method weren’t part of the problem, self.current would be sufficient on its own to handle popping the smallest element. Every time we call popSmallest(), we simply return self.current and then increment it to point to the next smallest number. This mechanism alone would allow us to simulate the behavior of an infinite set efficiently, without the need for storing all the integers in memory.

Example Without addBack():

  • Initially, self.current = 1.
  • Call popSmallest(): return 1 and set self.current = 2.
  • Call popSmallest() again: return 2 and set self.current = 3.
  • This pattern continues indefinitely, with self.current always returning the smallest available number.

In fact, if we didn’t need to reintroduce elements into the set, this mechanism would completely solve the problem. The design would be simple and efficient, requiring no additional data structures to manage.

Why Does self.current Only Increase?

The reason self.current only increases is to ensure that we never “go back” in the sequence of integers. Once a number is popped, we don’t need to keep track of it unless it’s explicitly added back through the addBack() method. This decision ensures the solution remains efficient by avoiding unnecessary checks or rollbacks in the sequence.

However, the problem does require handling situations where numbers are added back into the set. For that, we need additional mechanisms, which we’ll explore next.

Handling Reintroduced Elements with a Min-Heap

Since the problem requires adding numbers back to the set, we need a way to manage those elements. This is where the min-heap (self.added) comes into play. The heap ensures that whenever elements are added back, we can still efficiently pop the smallest one next.

  • Why a Min-Heap?
    The heap allows us to efficiently retrieve the smallest element among those that have been reintroduced. Each time an element is added back, it is pushed into the heap, and when we need the smallest element, we pop it from the heap in O(log n) time.

Example with addBack():

  • Initially, self.current = 4 (meaning 1, 2, 3 have already been popped).
  • We call addBack(2) to reintroduce 2.
  • Now, when we call popSmallest(), we first check the heap (self.added), and since 2 is smaller than 4, we return 2.

Without this heap structure, managing the reintroduced elements in the correct order would be inefficient.

Preventing Redundant Additions with self.addedSet

To avoid adding the same element back into the heap multiple times, we maintain a set (self.addedSet). This set tracks which elements have been added back, providing O(1) lookup time to ensure duplicates aren’t added.

Avoiding Unnecessary addBack() Calls

Another important constraint in the addBack() method is the condition that prevents adding numbers greater than or equal to self.current. This is because numbers greater than or equal to self.current haven’t been popped yet, so there’s no need to add them back.

Example:

  • If self.current = 5, and we call addBack(6), there’s no need to reintroduce 6 because it hasn’t been removed from the set yet.

This keeps the solution efficient and avoids unnecessary operations.

Popping the Smallest Element with popSmallest()

The popSmallest() method intelligently combines both self.current and the min-heap to determine the smallest element to return:

  • If the heap is empty, we return self.current.
  • If the heap contains elements, we return the smallest element from the heap (which would have been reintroduced earlier).

This ensures that both the natural sequence of numbers (self.current) and any reintroduced elements are handled in the correct order.

The SmallestInfiniteSet class combines a few simple yet effective strategies to manage an infinite set of integers. self.current alone is sufficient to simulate popping elements in order, but the problem requires reintroducing elements, which is handled efficiently using a min-heap and a set. These structures work together to ensure that both the natural sequence of numbers and any reintroduced numbers are popped in the correct order, while keeping operations efficient.

The key takeaway is that by incrementing self.current and using the heap for reintroduced elements, we avoid the need to store the entire set, making this a highly scalable solution.