Fast Python implementation with heap. Explained!


  • 0
    L

    Implementing heap is easy but it may got slowed down by a naive de-duplication algorithm, for example, we can simply keep popping the top elements in the heap until it ran out of duplicates. Here is how I accelerated things up a little bit.
    Notice that every number in the ugly queue is a product of the given primes, and duplicates are introduced when two prime numbers got flipped in the multiplication, so we can get rid of them by enforcing the ordering of numbers in the multiplication. To do this, I coupled each number in the ugly queue with the last prime number it's generated from, and I skip adding a new number to the heap if a smaller prime number was about to be multiplied with the old number coupled with a bigger prime.
    The heap stores the next smallest candidate for each given prime number and the index of the preceding ugly number the candidate is generated from. The prime number itself can be retrieved by a simple division.

    class Solution(object):
        def nthSuperUglyNumber(self, n, primes):
            """
            :type n: int
            :type primes: List[int]
            :rtype: int
            """
            ugly = [[0,0]] * n #ugly[i][0] = preceding ugly number * ugly[i][1]
            ugly[0] = [1,1]
            heap = zip(primes,[0]*len(primes)) #heap[i][0] = latent prime number * ugly[heap[i][1]][0]
            heapq.heapify(heap)
            
            for i in xrange(1,n):
                heapMin, fromIndex = heapq.heappop(heap)
                p = heapMin / ugly[fromIndex][0] #get the ending prime number from the heap
                ugly[i] = [heapMin,p]            #and update the ugly queue
                fromIndex += 1
                
                while ugly[fromIndex][1] > p: 
                    fromIndex += 1
                #skip the ugly numbers ending with bigger primes
                #replacing ">" with "<" also works, but kinda counter-intuitive
                heapq.heappush(heap,(ugly[fromIndex][0]*p,fromIndex)) #update the heap
                
            return ugly[-1][0]
    

Log in to reply
 

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