Super Ugly Number

Problem: Super Ugly Number

The solution to this problem evolves from solution to ugly numbers. We need to manage multiple arrays to decide which primes are used to generate current ugly number. A heap is a great way to manage different ways because we always need the minimum answer to append.

Code in Python:

import heapq

class Solution(object):
    def nthSuperUglyNumber(self, n, primes):
        """
        :type n: int
        :type primes: List[int]
        :rtype: int
        """
        if n == 1: return 1
        ans, k = [1], 1
        heap = []

        for prime in primes:
            heapq.heappush(heap, (prime, prime, 0))

        ugly, prime, index = heapq.heappop(heap)
        while True:
            if k == n - 1: return ugly
            ans.append(ugly)
            k += 1
            heapq.heappush(heap, (prime*ans[index+1], prime, index+1))
            _ugly, _prime, _index = heapq.heappop(heap)
            while ugly == _ugly:
                heapq.heappush(heap, (_prime*ans[_index+1], _prime, _index+1))
                _ugly, _prime, _index = heapq.heappop(heap)
            ugly, prime, index = _ugly, _prime, _index

However, this solution consumes a lot of space and time on push, top and append. A better way is to use generators in python, which can be seen from a solution here.