My priority_queue based C++ code, 276ms


  • 2
    D

    The straightforward solution is to use a big heap (i.e. priority_queue in C++) to store all currently generated ugly numbers and keep track of the minimum one. Everytime, we remove the minimum one and add the newly generated ugly numbers in the queue (i.e the minimum one X primes[i], for i=[0,K-1] and remove the duplicates). This will require roughly O(NK) space and O(Nlog(NK), worst case) time.
    It is very slow since there are many ordering information we didn't use in the above algorithm (e.g. minprime[i]< minprime[j], if i<j). So we should try to use such information to speed up our algorithm. Can we regroup the currently generated ugly number in a way that each group is in a ascending order and we only need to use a heap to find the min of the first elements of all the groups?

    One option would be divide all the currently generated ugly numbers in K groups, each time, we remove the top (i.e. min) number from priority_queue, we put the newly generated ugly number top()*primes[i] to the i-th group so that the i-th group includes all the currently generated ugly numbers that are multiple of primes[i]. Of course, this will include some duplicates (e.g 6 will be in the group of 2 and also in the group of 3). Since in this way, all the groups are in the ascending order (i.e. Group[i][m] > Group[i][n], if m>n). So we only need to use a K-element heap to track the next minimum ugly number. So the time complexity goes down. However, the memory complexity is still O(NK). We will get MLE error.

    Can we do better? One thing we notice is that for all elements in the i_th group has the format Ugly[m]*primes[i]. So we don't need to save all those elements but just Ugly[0..n-1]. The i_th group (i.e. multiple of primes[i]) has elements [Ugly[0] * primes[i], Ugly[1] * primes[i], ..., Ugly[m] * primes[i]]. We just need to track the index of Ugly (which indicates which Ugly element) should be used to generate the next minimum element in the i_th group, which will be push in the priority_queue. Now we reduce the memory complexity to O(N). But due to the duplicates issue, the time complexity could be O(NKlogK). Can we do further optimization?
    The below code has run time of 1020ms

    class Solution {
    public:
        int nthSuperUglyNumber(int n, vector<int>& primes) { // assuming inputs are valid
            int K = primes.size(), res[n], lastMul[n] ={0}, i, curVal, curIdx; //lastMul is the indices of K groups   
            priority_queue<pair<int, int>> headHeap; // K-element heap to find the minimum of the minimum of K groups
            for(res[0]=-1, i=0; i<K;++i) // initialize the heap
                headHeap.push(make_pair(-primes[i], i)); // first the ugly number value (change to negative to get the minimum one), second is the index of ugly number queues, i.e. candidates[i]
    
            for(i=0;i<n;)  //generate the first n ugly number
            { // the elements in the priority_queue are negative to help us find the minimum
                curVal = headHeap.top().first;
                curIdx = headHeap.top().second; // group index of the current minum ugly number
                headHeap.pop();
                if(curVal!=res[i]) res[++i] = curVal; // if it is not a duplicate, the top one will the ++i_th ugly number
                headHeap.push(make_pair(res[++lastMul[curIdx]] * primes[curIdx], curIdx)); // generate the next minimum element of the curIdx_th group and push it in the heap
            }
            return -res[n-1];
        }
    };
    

    (Thanks for StefanPochmann's comments that inspired this optimization)
    To avoid duplicates, we can use extra memory (i.e.pIdx) to save the index of the maximum prime factor of each ugly number that treat each ugly number only as multiple of its maximum prime factor. For example 6 is treated only as multiple of 3 , not multiple of 2. So it can only be generated as 2 x 3, not 3x2. In that way, no duplicate will occur. The below version has N heap insert operations and run time of 230ms.

    class Solution {
    public:
        int nthSuperUglyNumber(int n, vector<int>& primes) {
            int K = primes.size(), res[n], pIdx[n], lastMul[n] ={}, i, curVal, curIdx;
    
            priority_queue<pair<int, int>> headHeap; // 
            for(res[0]=-1, i=0; i<K;++i)   headHeap.push(make_pair(-primes[i], i)); // first the ugly number value (change to negative to get the minimum one), second is the index of ugly number queues, i.e. candidates[i]
    
            for(i=1;i<n;++i)
            {
                curVal = headHeap.top().first;
                curIdx = headHeap.top().second;
                headHeap.pop();
                res[i] = curVal; pIdx[i] = curIdx; // pIdx is the index of the maximum prime factor
                while(pIdx[++lastMul[curIdx]]>curIdx); // skip ugly numbers whose maximum factor index is larger than curIdx.
                headHeap.push(make_pair(res[lastMul[curIdx]] * primes[curIdx], curIdx));
            }
            return -res[n-1];
        }
    };

  • 0

    I don't see why it's O(NlogK). Each loop iteration costs O(logK) and you do more than N iterations (because of the duplicates).


  • 0
    D

    You are right, worst case is O(NKlogk), but if we can use more memory, we can avoid such duplicates issue. So I added another version, which should be O(NlogK) since duplicates are avoided. Your further comments/suggestions are welcome


  • 0
    X

    At the end, lastMul for each of the prime number will be close to N (Or the size increase linearly with N).
    It will be O(NK).


  • 0

    As @xiaohl0913 said, At the end, lastMul for each of the prime number will be close to N (Or the size increase linearly with N). The overall time complexity will be O(N^2).


  • 0
    D

    You are right that O(NlogK) complexity is not accurate. I just considered the heap operation complexity as the algorithm complexity, which is not accurate. O(NK) is a upper bound. Will think about how to further optimize it to achieve O(nlogk). Thanks for your comments..


  • 0
    D

    Thanks for your comments. I guess you mean O(NK), which I agree is a upper bound. I also agree O(NlogK) is not accurate. But O(N^2) is wrong.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.