2336.smallest-number-in-infinite-set.md

2336. Smallest Number in Infinite Set

Test Cases

Edge / Corner Cases

• Add the number which is still in the infinite set.

1, 2, 3, 4, 5, 6
   s
         * add(4)

• Duplicate adding.

1, 2, 3, 4, 5, 6
            s
   * add(2) // Added back successfully
   * add(2) // Should be ignored

Heap + Set (or TreeSet)

We can add all possible numbers to the set as the initial set, and use a minimum heap to find the smallest number in the set, including the numbers that are added back. But this approach is not efficient, because the set could be infinite.

How can we solve this problem efficiently? Let's break down the problem:

• If we don't add number back, we can just use a variable smallest to store the current smallest number in the infinite set.

1, 2, 3, 4, 5, 6, ...
      * ->
      smallest

• If we support adding numbers back, we need a data structure that can find the smallest number dynamically, which is heap. And we need a set to avoid duplicate adding. So if we add the number back, we can add it to the heap and set.

1, 2, 3, 4, 5, 6, ...
|-----------|  * ->
    heap       smallest

To summarize to our approach:

• popSmallest(): we check if heap is empty, if not, there are some numbers added back, we just pop the smallest number from the heap. If the heap is empty, we return the current smallest number in the infinite set, that is maintained by a variable.
• addBack(number): we check if the number is still in the infinite set, if not, we add it to the heap and set.

class SmallestInfiniteSet() {
    // The smallest number in the infinite set.
    private var smallest = 1

    // The heap and set to store the numbers that are added back.
    private val minHeap = PriorityQueue<Int>()
    private val added = HashSet<Int>()

    fun popSmallest(): Int {
        // If we have added some numbers back, we just pop the smallest number from the heap.
        if (heap.isNotEmpty()) {
            val value = minHeap.poll()
            added.remove(value)
            return value
        } else {
            // Otherwise, we just return the current smallest number in the infinite set.
            return smallest++
        }
    }

    fun addBack(number: Int) {
        if (smallest <= number ||   // If we add the number that is still in the infinite set.
            number in set) {        // Or we add the number that is already added back.
            return
        }

        added.add(number)
        minHeap.add(number)
    }
}

// Or equivalently using TreeSet
class SmallestInfiniteSet() {
    private val addBack = TreeSet<Int>() // To store the numbers that are removed and added back
    private var smallest = 1

    fun popSmallest(): Int {
        if (addBack.isNotEmpty()) {
            val value = addBack.pollFirst()
            return value
        } else {
            return smallest++
        }
    }

    // Same as the previous approach
    fun addBack(number: Int) {
        if (smallest <= number || addBack.contains(number)) return
        addBack.add(number)
    }
}

• Time Complexity: O(log n) for popSmallest(), O(log n) for addBack()
• Space Complexity: O(n)