Leetcode 313. Super Ugly Number

Problem Explanation

This problem asks us to find the nth Super Ugly Number given a list of primes.

Super Ugly Numbers are positive numbers whose all prime factors are in the given prime list primes of size k. For example, if the primes list is [2,7,13,19], the first 12 Super Ugly Numbers would be [1,2,4,7,8,13,14,16,19,26,28,32]. So for n=12, the output would be 32.

Approach

The solution is a modification of the regular Ugly Number II problem which only considers the primes [2,3,5]. The idea is to initialize a list of ugly numbers with 1 and maintain a list of indices, one for each prime. Each index indicates the smallest ugly number that has not been multiplied by the corresponding prime. We make multiple iterations, each time, the smallest number among the product of the ugly number at the indices and the corresponding prime is the next ugly number. We then increase the indices of all primes that divide this smallest number.

For example, if the primes list is [2,7,13,19], initially, the list of ugly numbers is [1] and the list of indices is [0,0,0,0] corresponding to each prime. In the first iteration, the smallest number among [2*1, 7*1, 13*1, 19*1] is 2. Hence, 2 is added to the list of ugly numbers and the index of 2 is increased by 1. Now, the list of ugly numbers is [1,2] and indices are [1,0,0,0]. We continue these steps until we have n ugly numbers in the list.

This approach ensures that we generate all super ugly numbers in ascending order without any duplicates.

To implement our solution, we simply need to understand the problem and apply standard data structures like arrays, vectors and a priority queue available in all popular programming languages.

With this understanding, let's look at how to implement this in different languages.

Python Solution

1
2python
3import heapq
4class Solution:
5    def nthSuperUglyNumber(self, n: int, primes: List[int]) -> int:
6        ugly = [1]
7        heap = [(prime, prime, 0) for prime in primes]
8        while n > 1:
9            val, prime, idx = heapq.heappop(heap)
10            if val > ugly[-1]:
11                n -= 1
12                ugly.append(val)
13            heapq.heappush(heap, (prime * ugly[idx+1], prime, idx+1))
14        return ugly[-1]

C++ Solution

1
2cpp
3#include <algorithm>
4#include <vector>
5using namespace std;
6class Solution {
7public:
8    int nthSuperUglyNumber(int n, vector<int>& primes) {
9        const int k = primes.size();
10        vector<int> indices(k);
11        vector<int> uglyNums{1};
12
13        while (uglyNums.size() < n) {
14            vector<int> nexts(k);
15            for (int i = 0; i < k; ++i)
16                nexts[i] = uglyNums[indices[i]] * primes[i];
17            const int next = *min_element(begin(nexts), end(nexts));
18            for (int i = 0; i < k; ++i)
19                if (next == nexts[i])
20                    ++indices[i];
21            uglyNums.push_back(next);
22        }
23
24        return uglyNums.back();
25    }
26};

Java Solution

1
2java
3import java.util.PriorityQueue;
4
5class Solution {
6    public int nthSuperUglyNumber(int n, int[] primes) {
7        PriorityQueue<long[]> pq = new PriorityQueue<>((a, b) -> Long.compare(a[0], b[0]));
8        for (int prime : primes) pq.add(new long[]{prime, prime, 1});
9        int[] ugly = new int[n];
10        ugly[0] = 1;
11        for (int i = 1; i < n; i++) {
12            long[] poll = pq.poll();
13            ugly[i] = (int)poll[0];
14            while (!pq.isEmpty() && pq.peek()[0] == ugly[i]) poll = pq.poll();
15            pq.add(new long[]{poll[1] * ugly[(int)poll[2]], poll[1], poll[2] + 1});
16        }
17        return ugly[n - 1];
18    }
19}

JavaScript Solution

1
2javascript
3var nthSuperUglyNumber = function (n, primes) {
4    let k = primes.length,
5        next = Array(k).fill(0),
6        uglyNums = [1];
7
8    while (uglyNums.length < n) {
9        let min = Math.min(...(primes.map((val, idx) => val * uglyNums[next[idx]])));
10        if (min != uglyNums[uglyNums.length - 1]) {
11            uglyNums.push(min);
12        }
13        next = next.map((val, idx) => (primes[idx] * uglyNums[val] == min ? val + 1 : val));
14    }
15
16    return uglyNums[n - 1];
17};

C# Solution

1
2csharp
3public class Solution {
4    public int NthSuperUglyNumber(int n, int[] primes) {
5        int[] ugly = new int[n];
6        int[] primeIndices = new int[primes.Length];
7        ugly[0] = 1;
8        for (int idx = 1; idx < n; idx++) {
9            ugly[idx] = int.MaxValue;
10            for (int j= 0; j < primes.Length; j++) {
11                ugly[idx] = Math.Min(ugly[primeIndices[j]] * primes[j], ugly[idx]);
12            }
13            for (int j= 0; j < primes.Length; j++) {
14                if (ugly[idx] == ugly[primeIndices[j]] * primes[j]) primeIndices[j]++;
15            }
16        }
17        return ugly[n - 1];
18    }
19}

Rust Solution

1
2rust
3pub fn nth_super_ugly_number(n: i32, primes: Vec<i32>) -> i32 {
4    let mut ugly = vec![1];
5    let mut idx = vec![0; primes.len()];
6
7    for _ in 1..n {
8        ugly.push(
9            (0..primes.len()).map(|i| ugly[idx[i]] * primes[i])
10                .min().unwrap()
11        );
12        for i in 0..primes.len() {
13            if ugly[idx[i]] * primes[i] == *ugly.last().unwrap() {
14                idx[i] += 1;
15            }
16        }
17    }
18
19    *ugly.last().unwrap()
20}

The above Rust solution was achieved by creating a vector of super ugly numbers "ugly" and a vector "idx" to keep track of the smallest ugly number that has not been multiplied by the corresponding prime. Just like in the other solutions, it pushes the minimum product of the ugly number and the corresponding prime in each iteration, advancing the indices of all primes that divide this minimum value.

Perl Solution

1
2perl
3sub nth_super_ugly_number {
4    my ($n, $primes) = @_;
5    my @ugly = (1);
6    my @idx = map { 0 } @$primes;
7
8    for (2..$n) {
9        my @next = map { $ugly[$idx[$_]] * $primes->[$_] } 0..$#$primes;
10        my $min = List::Util::min(@next);
11        push @ugly, $min;
12        for (0..$#$primes) {
13            $idx[$_]++ if $next[$_] == $min;
14        }
15    }
16
17    return $ugly[-1];
18}

This Perl solution implements the same logic with an array @ugly for super ugly numbers and @idx index array. The List::Util::min function is used to get the minimum value from the @next array. The index of all primes that divide this minimum value is increased. This solution requires the Perl module List::Util.

Swift Solution

1
2swift
3class Solution {
4    func nthSuperUglyNumber(_ n: Int, _ primes: [Int]) -> Int {
5        var ugly = [Int](repeating: 0, count: n)
6        ugly[0] = 1
7        var index = [Int](repeating: 0, count: primes.count)
8
9        for i in 1..<n {
10            ugly[i] = Int.max
11            for j in 0..<primes.count {
12                ugly[i] = min(ugly[i], primes[j] * ugly[index[j]])
13            }
14            for j in 0..<primes.count {
15                if ugly[i] == primes[j] * ugly[index[j]] {
16                    index[j] += 1
17                }
18            }
19        }
20        return ugly[n - 1]
21    }
22}

In the Swift solution, min function is used to find the smallest super ugly number - the minimum product of the ugly number and the prime. The index of all primes that divide this minimum value is increased.