[LeetCode] 2336. Smallest Number in Infinite Set

1. Question

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.

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
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.

2. Thoughts

1. Tried to think of a scheme that would work in real infinite time.
2. Yet couldn't and because number of operation is limited to 1000, set 1000 elements in a heap, and calculated the min

3. Tips learned

4. My solution

class SmallestInfiniteSet:
    def __init__(self):
        self.infiniteHeap = []
        self.infiniteSet = set()
        for i in range(1, 1001):
            heapq.heappush(self.infiniteHeap, i)
            self.infiniteSet.add(i)

    def popSmallest(self) -> int:
        smallest = heapq.heappop(self.infiniteHeap)
        self.infiniteSet.remove(smallest)
        return smallest

    def addBack(self, num: int) -> None:
        if num not in self.infiniteSet:
            heapq.heappush(self.infiniteHeap, num)
            self.infiniteSet.add(num)

137/137 cases passed (103 ms)
Your runtime beats 64.41 % of python3 submissions
Your memory usage beats 57.67 % of python3 submissions (17.3 MB)

5. AI Solution and Improvements

Your implementation of the SmallestInfiniteSet class is quite efficient and uses a min-heap (via Python's heapq) and a set to manage the integers. This approach allows for fast access to the smallest number and quick checks for the existence of a number in the set. However, there are some areas for improvement and optimization:

1. Infinite Set Initialization: Your current approach initializes the heap and set with the first 1000 positive integers. This is a limitation because the set is supposed to be infinite. We can avoid pre-populating a finite number of integers and instead generate them on demand.

2. Populating the Heap: When popping the smallest element, if the heap becomes empty, we need to add the next integer to it. This can be managed by keeping track of the next integer to be added.

3. Add Back Optimization: Your current implementation checks if the number is not in the set before adding it back. This is correct, but we can optimize it further.

4. Space Complexity: Using both a heap and a set increases space complexity. We can optimize it by using only a set to track removed numbers and generating other numbers on demand.

Here's an optimized version of your class:

import heapq

class SmallestInfiniteSet:
    def __init__(self):
        self.removed = set()
        self.next_num = 1  # The next number to be added if the heap is empty

    def popSmallest(self) -> int:
        # If there are removed elements in the heap, pop the smallest
        if self.removed:
            smallest = heapq.heappop(self.removed)
        else:
            # Else, use the next available integer and increment it
            smallest = self.next_num
            self.next_num += 1
        return smallest

    def addBack(self, num: int) -> None:
        # Add the number back only if it's smaller than the next number and not already in the heap
        if num < self.next_num and num not in self.removed:
            heapq.heappush(self.removed, num)

This implementation handles an infinite set of positive integers more effectively by generating them on demand and managing only the removed numbers in a heap.