Sharing very simple and elegant Python solution using heap, with explanation


  • 11
    import heapq
    
    class Solution(object):
        
        # The idea is to keep a heap that stores a tuple (val, fact).
        # The fact value stores the last factor the number was multiplied to.
        # So every time I extract the minimum and look at the queue, I know
        # it must be multiplied by queue and next factor, otherwise I will have a
        # duplicate. 
        #
        # For example, we have 2 and 5 with fact 2 and fact 5. When I extract
        # 2 and fact 2, I will insert 2 * 5 = 10 with fact 5. Then, when I extract 5 from queue
        # 5, if I do 5 * 2 (fact 2) = 10,
        # I will have a duplicate. So every number coming from
        # the k-th fact must be multiplied only by the factor fact and higher.
        
        
        
        def nthUglyNumber(self, n):
            if n is 1:
                return 1
            h, val = [(2, 2), (3, 3), (5, 5)], 1
            for i in range(1, n):
                val, fact = heapq.heappop(h)
                heapq.heappush(h, (val * 5, 5))
                if fact <= 3:
                    heapq.heappush(h, (val * 3, 3))
                if fact is 2:
                    heapq.heappush(h, (val * 2, 2))
            return val

  • 9

    Gets simpler with a more basic base case...

    def nthUglyNumber(self, n):
        h = [(1, 1)]
        for _ in range(n):
            val, fact = heapq.heappop(h)
            heapq.heappush(h, (val * 5, 5))
            if fact <= 3:
                heapq.heappush(h, (val * 3, 3))
            if fact <= 2:
                heapq.heappush(h, (val * 2, 2))
        return val
    

    And after unifying the cases:

    def nthUglyNumber(self, n):
        h = [(1, 1)]
        for _ in range(n):
            val, fact = heapq.heappop(h)
            for x in 2, 3, 5:
                if fact <= x:
                    heapq.heappush(h, (val * x, x))
        return val

  • 0

    Right... Thank you Stefan.


  • 0
    E

    Thanks @agave and @StefanPochmann , I have rewritten it into cpp version : )

    int nthUglyNumber(int n) {
        priority_queue< pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> pq;
        pq.push(make_pair(1,1));
        int val, fact;
        for(int i = 0; i < n ; i++){
            auto t = pq.top(); 
            val = t.first; 
            fact = t.second;
            pq.pop();
            // handle overflow problem
            if( fact <= 2 && val < 1073741823.5 ) pq.push(make_pair(val*2, 2));
            if( fact <= 3 && val <  715827882.3 ) pq.push(make_pair(val*3, 3));
            if( fact <= 5 && val <  429496729.4 ) pq.push(make_pair(val*5, 5));
        }
        return val;
    }
    

Log in to reply
 

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