LeetCode 313: Super Ugly Number

Problem Restatement

A super ugly number is a positive integer whose prime factors all come from the array primes.

We are given:

n
primes

Return the nth super ugly number.

The sequence starts with 1. The number 1 is included even though it has no prime factors.

The official statement says the nth super ugly number is guaranteed to fit in a 32-bit signed integer. The prime list is unique and sorted in ascending order.

Input and Output

Function shape:

def nthSuperUglyNumber(n: int, primes: list[int]) -> int:
    ...

Examples

Example 1:

n = 12
primes = [2, 7, 13, 19]

The first 12 super ugly numbers are:

[1, 2, 4, 7, 8, 13, 14, 16, 19, 26, 28, 32]

Output:

32

Example 2:

n = 1
primes = [2, 3, 5]

Output:

1

The first super ugly number is always 1.

First Thought: Generate Every Number

One direct idea is to check every positive integer:

1. Factor the number.
2. Check whether every prime factor belongs to primes.
3. Count valid numbers until reaching the nth one.

This works in theory.

But it wastes time on numbers that are not super ugly.

For large n, this is far too slow.

Key Insight

Every super ugly number after 1 is made by multiplying an earlier super ugly number by one of the allowed primes.

So if we already know:

ugly = [1, ...]

then future candidates look like:

ugly[i] * prime

for some earlier index i and some prime.

This is similar to merging several sorted streams.

For each prime p, imagine a stream:

p * ugly[0], p * ugly[1], p * ugly[2], ...

The next super ugly number is the smallest value among all stream heads.

DP State

Let:

ugly[i]

be the (i + 1)th super ugly number.

So:

ugly[0] = 1

For each prime primes[j], keep a pointer:

idx[j]

This pointer tells us which existing ugly number should be multiplied by primes[j].

The current candidate for prime j is:

primes[j] * ugly[idx[j]]

At every step:

1. Take the smallest candidate.
2. Append it to ugly.
3. Advance every pointer whose candidate produced that value.

The third step is important because duplicates can happen.

For example, with primes [2, 3]:

6 = 2 * 3 = 3 * 2

We should include 6 only once.

Algorithm

Initialize:

ugly = [1] * n
idx = [0] * len(primes)

Then for every position from 1 to n - 1:

1. Compute all current candidates.
2. Let next_value be the minimum candidate.
3. Store it in ugly[i].
4. For every prime whose candidate equals next_value, advance its pointer.

Finally return:

ugly[n - 1]

Correctness

The algorithm starts with 1, which is the first super ugly number.

For each prime, the pointer idx[j] always points to the smallest known ugly number that can still produce a future candidate with primes[j].

At each step, every possible next super ugly number must be formed as:

previous_super_ugly * allowed_prime

The algorithm keeps the smallest available product for every allowed prime. Therefore, the smallest value among these candidates is the next smallest super ugly number.

When multiple primes produce the same value, advancing all matching pointers prevents duplicates while preserving all future candidates.

By induction, after filling ugly[i], the array contains the first i + 1 super ugly numbers in sorted order.

Therefore, ugly[n - 1] is the required nth super ugly number.

Complexity

Let:

This dynamic programming method avoids heap duplicates and is efficient for the usual constraints.

Implementation

class Solution:
    def nthSuperUglyNumber(self, n: int, primes: list[int]) -> int:
        k = len(primes)

        ugly = [1] * n
        idx = [0] * k

        for i in range(1, n):
            candidates = [
                primes[j] * ugly[idx[j]]
                for j in range(k)
            ]

            next_value = min(candidates)
            ugly[i] = next_value

            for j in range(k):
                if candidates[j] == next_value:
                    idx[j] += 1

        return ugly[n - 1]

Code Explanation

We store the number of primes:

k = len(primes)

The DP array starts with 1.

ugly = [1] * n

The pointer array starts at zero.

idx = [0] * k

At each position, we compute one candidate per prime:

candidates = [
    primes[j] * ugly[idx[j]]
    for j in range(k)
]

The smallest candidate is the next super ugly number.

next_value = min(candidates)
ugly[i] = next_value

Then we advance every pointer that produced this value.

for j in range(k):
    if candidates[j] == next_value:
        idx[j] += 1

This removes duplicates.

For example, if both 2 * 3 and 3 * 2 produce 6, both pointers move forward, and 6 appears only once.

Finally:

return ugly[n - 1]

Because ugly[0] is the first super ugly number.

Slightly Faster Version

The previous implementation recomputes the candidate list every loop.

We can keep candidates in an array and update only changed entries.

class Solution:
    def nthSuperUglyNumber(self, n: int, primes: list[int]) -> int:
        k = len(primes)

        ugly = [1] * n
        idx = [0] * k
        candidates = primes[:]

        for i in range(1, n):
            next_value = min(candidates)
            ugly[i] = next_value

            for j in range(k):
                if candidates[j] == next_value:
                    idx[j] += 1
                    candidates[j] = primes[j] * ugly[idx[j]]

        return ugly[-1]

This keeps the same asymptotic complexity, but reduces repeated multiplication.

Testing

def run_tests():
    s = Solution()

    assert s.nthSuperUglyNumber(12, [2, 7, 13, 19]) == 32
    assert s.nthSuperUglyNumber(1, [2, 3, 5]) == 1
    assert s.nthSuperUglyNumber(10, [2, 3, 5]) == 12
    assert s.nthSuperUglyNumber(5, [7]) == 2401
    assert s.nthSuperUglyNumber(6, [2]) == 32
    assert s.nthSuperUglyNumber(8, [3, 5]) == 45

    print("all tests passed")

run_tests()

Test meaning: