313_SuperUglyNumber.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# @Author: xuezaigds@gmail.com
# @Last Modified time: 2016-07-10 10:32:18


class Solution(object):
    """ Basic idea is same as ugly number II.

    New ugly number is generated by multiplying a prime with
    previous generated ugly number.
    One catch is need to remove duplicate.

    Make code shorter and faster according to:
    https://discuss.leetcode.com/topic/33103/fast-python-solution-based-on-solution-for-ugly-number-ii
    """
    def nthSuperUglyNumber(self, n, primes):
        k = len(primes)
        ugly_numbers = [1] * n
        index_record = [-1] * k
        value_record = [1] * k

        ugly_index = 0
        while ugly_index < n:
            min_val = min(value_record)
            ugly_numbers[ugly_index] = min_val

            for i in range(k):
                # Remove duplicate numbers here.
                if value_record[i] == min_val:
                    index_record[i] += 1
                    value_record[i] = primes[i] * ugly_numbers[index_record[i]]

            ugly_index += 1

        return ugly_numbers[n - 1]

"""
2
[2,3,5]
100
[2,3,5,7,11]
5
[2,3,5,7,11,13,17,19]
"""