Smallest Number in Infinite Set

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

class SmallestInfiniteSet:
    import heapq
    def __init__(self):
        self.summary = list(range(1, 1001))
        heapq.heapify(self.summary)
        self.deleted = set()
    def popSmallest(self) -> int:
        A = heapq.heappop(self.summary)
        self.deleted.add(A)
        return A
    def addBack(self, num: int) -> None:
        if num in self.deleted:
            self.deleted.remove(num)
            heapq.heappush(self.summary, num)