Could anyone help with my TLE with Python using Min-heap?


  • 0
    S

    The idea is similar to merge k sorted arrays, and also similar to Ugly number2. I used a minheap for extract the smallest result each time, but it still TLE.

    def nthSuperUglyNumber(self, n, primes):
        """
        :type n: int
        :type primes: List[int]
        :rtype: int
        """
        if n == 1:
            return 1
        res = [1]
        heap = []
        k = len(primes)
        for i in xrange(k):
            heapq.heappush(heap,(primes[i],0,primes[i]))
    
        while len(res) < n:
            nxt, index, prime = heap[0]
            res+=[nxt]
            while heap and heap[0][0] == nxt:
                heapq.heappop(heap)
                heapq.heappush(heap, (prime*res[index+1],index+1,prime))
        return res[-1]

Log in to reply
 

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