C++     Design     Hash Table     Heap (Priority Queue)     Medium     Ordered Set    

Problem Statement:

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.

Solution:

My first idea was to use a min heap and store the elements which are added.

 
class SmallestInfiniteSet {
public:
    priority_queue<int,vector<int>,greater<int>> pq;
    int bot = 1;
    SmallestInfiniteSet() 
    {}
    
    int popSmallest() 
    {
        if (pq.empty()) return bot++;
        int res = pq.top();
        pq.pop();
        return res;
    }
    
    void addBack(int num) 
    {
        if(num<bot) pq.push(num);
    }
};
 

This solution works for 90% of tests but will fail when a number is added back twice, hence it will be present in the heap twice. So, we can track this using a separate hashset.

 
class SmallestInfiniteSet {
public:
    priority_queue<int,vector<int>,greater<int>> pq;
    unordered_set<int> myset;
    int bot = 1;
    SmallestInfiniteSet() {}
    
    int popSmallest() 
    {
        if (pq.empty()) return bot++;
        int res = pq.top();
        pq.pop();
        myset.erase(res);
        return res;
    }
    
    void addBack(int num) 
    {
        if (!(num<bot && !myset.count(num))) return;
        pq.push(num);
        myset.insert(num);
    }
};
 

This solution gives us AC and beats 95% of all solutions in runtime.

At this point I noticed that the question has constrained itself to [1,1000] so we can exploit this to create a much easier solution

 
class SmallestInfiniteSet {
public:
    set<int> myset;
    SmallestInfiniteSet() 
    {
        for(int i=1; i<=1000; i++) myset.insert(i);
    }
    
    int popSmallest() 
    {
        int res = *myset.begin();
        myset.erase(res);
        return res;
    }
    
    void addBack(int num) 
    {
        myset.insert(num);
    }
};
 

This also gives us an AC but beats only 40% of all solutions in runtime.

Analysis: Both solutions have $O((m+n)\log(n))$ TC and $O(n)$ SC. Here $m$ is the number of calls and $n$ is maximum size of num.