Problem statement

https://leetcode.com/problems/super-ugly-number/

Solution 1

First solution is extension of problem 0264, but now we need to keep not 3 pointers, but k of them.

Complexity

Time complexity is O(nk), space is O(k).

Code

class Solution:
    def nthSuperUglyNumber(self, n, primes):
        k = len(primes)
        starts, nums = [0] * k, [1]
        for i in range(n-1):
            candidates = [primes[i]*nums[starts[i]] for i in range(k)]
            new_num = min(candidates)
            nums.append(new_num)
            starts = [starts[i] + (candidates[i] == new_num) for i in range(k)]
        return nums[-1]

Solution 2

Previous solution can improved to O(n log k), if we use heap to maintain all indexes in our heap with tuples (p*u, p, i), where p*u is the smallest value divisible by p, such that it is not traversed yet. We can look at it at merge of k lists, where each time we choose the smallest element.

Complexity

Time complexity is O(n log k).

Code

class Solution:
    def nthSuperUglyNumber(self, n, primes):
        cand = [(p, p, 1) for p in primes]
        ugly = [1]
        for _ in range(n-1):
            ugly.append(cand[0][0])
            while cand[0][0] == ugly[-1]:
                x, p, i = heappop(cand)
                heappush(cand, (p*ugly[i], p, i+1))
        return ugly[-1]