Super Ugly Number

题目描述

Write a program to find the nth super ugly number.

Super ugly numbers are positive numbers whose all prime factors are in the given prime list primes of size k. For example, [1, 2, 4, 7, 8, 13, 14, 16, 19, 26, 28, 32] is the sequence of the first 12 super ugly numbers given primes = [2, 7, 13, 19] of size 4.

解题方法

Solution

from heapq import *
class Solution(object):
    def nthSuperUglyNumber(self, n, primes):
        """
        :type n: int
        :type primes: List[int]
        :rtype: int
        """
        if not n:
            return 
        if n == 1:
            return 1

        heap = []
        ugly_nums = [1]
        length_p = len(primes)
        indexes = [0 for i in range(length_p)]
        for i in range(length_p):
            heappush(heap, primes[i] * ugly_nums[indexes[i]])

        for i in range(1, n):
            min_ugly_num = heap[0]
            ugly_nums.append(min_ugly_num)
            for i in range(length_p):
                if primes[i] * ugly_nums[indexes[i]] == min_ugly_num:
                    indexes[i] += 1
                    # pop and add next into heap
                    heappop(heap)
                    heappush(heap, primes[i] * ugly_nums[indexes[i]])

        return ugly_nums[n-1]

Reference