[leetcode] 2336. Smallest Number in Infinite Set _ Algorithm Problem Solve for python



1. Problem

2336. Smallest Number in Infinite Set

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.

Constraints:

  • 1 <= num <= 1000
  • At most 1000 calls will be made in total to popSmallest and addBack.

2. Solution

I solved this problem like this.

  • I made the heap myself!
class SmallestInfiniteSet:

    def __init__(self):
        self.heap = [None]
        self.heap_dic = {}
        for x in range(1, 1001):
            self.addBack(x)

    def popSmallest(self) -> int:
        smallest = self.heap[1]
        del self.heap_dic[smallest]

        self.heap[1] = self.heap[-1]
        del self.heap[-1]

        # make heap
        parent = 1
        child = parent * 2
        while (child < len(self.heap)):
            child1, child2 = child, child + 1

            smaller = child1
            if not (child2 >= len(self.heap)): # pick smaller children
                if self.heap[child1] > self.heap[child2]:
                    smaller = child2
            
            if self.heap[smaller] < self.heap[parent]:
                self.heap[parent], self.heap[smaller] = self.heap[smaller], self.heap[parent]

            parent = smaller
            child = smaller * 2

        return smallest


    def addBack(self, num: int) -> None:
        if num in self.heap_dic:
            return 

        self.heap.append(num)
        self.heap_dic[num] = True
        
        # make heap
        end = len(self.heap) - 1
        parent = end // 2
        child = end
        while (parent > 0) and (self.heap[parent] > self.heap[child]):
            self.heap[parent], self.heap[child] = self.heap[child], self.heap[parent]
            parent = parent // 2
            child = child // 2

        return

        
#    1
#  2   3
# 4 5 6 7

#         0
#    1          2
#  3   4     5     6
# 7 8 9 10 11 12 13 14

# Your SmallestInfiniteSet object will be instantiated and called as such:
# obj = SmallestInfiniteSet()
# param_1 = obj.popSmallest()
# obj.addBack(num)