Why my self-implemented heap get TLE ?


  • 0
    J

    When I use the heapq module of python I can pass the test:

    python

        import heapq
        hp = [1]
        i = 0
        prev = None
        dic = set()
        while i<n:
            nn = heapq.heappop(hp)
            if nn not in dic: 
                i += 1
                dic.add(nn)
                heapq.heappush(hp, nn*2)
                heapq.heappush(hp, nn*3)
                heapq.heappush(hp, nn*5)
        return nn
    

    But I implemented a heap, and get TLE for n = ~1600, can anybody see errors/defaults in my implementation ?

    class Solution:
    # @param {integer} n
    # @return {integer}
    def nthUglyNumber(self, n):
        heap = [-233] # heap array is indexing from 1 ! 
        
        def siftup():
            i = len(heap)-1
            while i>1 and heap[i]<heap[i/2]:
                i, heap[i], heap[i/2] = i/2, heap[i/2], heap[i]
                
        def siftdown():
            i = 1
            while i*2<len(heap):
                j = i
                if heap[i*2]<heap[i]: j = i*2 # j is the index of smallest element among the three
                if i*2+1<len(heap) and heap[i*2+1]<heap[j]: j = i*2+1
                if j!=i:
                    i, heap[i], heap[j] = j, heap[j], heap[i]
                else: break 
                
        def add(n):
            heap.append(n)
            siftup()
        
        def pop():
            n= heap[1]
            if len(heap)>2: 
                heap[1] = heap.pop()
                siftdown()
            else: heap.pop()# if only has one value in heap, after this pop, there will be nothing left
            return n
        
        add(1)
        i = 0
        prev = None
        while i<n:
            nn = pop()
            if nn != prev:
                i += 1
                prev = nn
                add(nn*2)
                add(nn*3)
                add(nn*5)
        return nn

Log in to reply
 

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