LeetCode 2336. Smallest Number in Infinite Set

Post by ailswan Jul. 11, 2024

2336. Smallest Number in Infinite Set

Problem Statement

link: LeetCode.cn LeetCode 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:

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

Solution Approach

The solution approach involves using a TreeSet to maintain a dynamically sorted set of added back elements and an integer to track the next smallest integer, ensuring efficient retrieval and reinsertion of the smallest element from the infinite set.

Algorithm

  1. Dynamic Ordering with TreeSet: Use a TreeSet to maintain a dynamically sorted collection of added back elements, allowing efficient retrieval and removal of the smallest element.
  2. Threshold Tracking: Maintain a threshold integer (thres) to keep track of the next smallest integer to be returned when the TreeSet is empty, ensuring a continuous sequence of positive integers.
  3. Efficient Operations: Implement popSmallest() to either return the smallest element from the TreeSet or the current threshold, and addBack() to add elements back into the TreeSet only if they are below the current threshold and not already present.

Implement

    class SmallestInfiniteSet: def __init__(self):
    self.cur = 1
    self.addlist = []
    self.addset = set()

def popSmallest(self) -> int:
    if self.addlist:
        ret = heapq.heappop(self.addlist)
        self.addset.remove(ret)
    else:
        ret = self.cur
        self.cur += 1
    return ret

def addBack(self, num: int) -> None:
    if num >= self.cur or num in self.addset:
        return
    heapq.heappush(self.addlist, num)
    self.addset.add(num)