313. Super Ugly Number

Description

Given an integer n and a list of prime numbers, primes, the task is to find the n-th super ugly number. A super ugly number is a positive integer whose only prime factors are listed in primes. The sequence of super ugly numbers is defined such that the i-th element is the smallest number whose prime factors are only from the given list.

Intuition

To determine the n-th super ugly number efficiently, we need a strategy that avoids generating all possible ugly numbers up to n, which would be infeasible for large values. Instead, we use a priority queue to dynamically track and generate the sequence of ugly numbers in a sorted order.

Approach

1. Priority Queue Initialization: - Use a min-heap (priority queue) to always access the smallest possible super ugly number efficiently. - Initialize the priority queue with the first prime multiples.
2. Generate Super Ugly Numbers: - Start with 1 (the first super ugly number) and iteratively generate new super ugly numbers by multiplying the smallest number in the priority queue with each prime from the primes list. - To avoid duplicates, maintain a list of generated numbers and ensure no duplicates are added.
3. Populate Output: - Extract the smallest number from the priority queue and add it to the output list. - Update the priority queue with new numbers generated by multiplying the extracted number by each prime factor.
4. Return Result: - Continue the process until the n-th super ugly number is found, which will be the result.

Complexity

Time Complexity:

• The time complexity is (O(n \log k)), where (k) is the number of primes. Each insertion and extraction operation from the priority queue takes (O(\log k)) time, and this is done n times.

Space Complexity:

• The space complexity is (O(n)) to store the sequence of super ugly numbers and the priority queue. The output array requires (O(n)) space, and the priority queue stores at most (O(k)) elements at a time, where (k) is the number of primes.

Code

C++

class Solution {
public:
    struct Compare {
        bool operator()(const vector& a, const vector& b) {
            return a[0] > b[0];
        }
    };

    int nthSuperUglyNumber(int n, vector& primes) {
        priority_queue, vector>, Compare> minPQ;

        // Initialize the priority queue with the first prime multiples
        for (int prime : primes) {
            minPQ.push({prime, 0, prime}); // {value, index in output, prime factor}
        }

        vector output(1, 1); // The first super ugly number is always 1

        while (output.size() < n) {
            vector current = minPQ.top();
            minPQ.pop();

            long value = current[0], idx = current[1], prime = current[2];

            // Avoid adding duplicate values to the output array
            if (value != output.back()) {
                output.push_back(value);
            }

            // Generate the next multiple of the current prime
            minPQ.push({output[idx + 1] * prime, idx + 1, prime});
        }

        return output[n - 1];
    }
};

Python

class Solution:
    def nthSuperUglyNumber(self, n: int, primes: list[int]) -> int:
        minPQ = []
        
        # Initialize the priority queue with the first prime multiples
        for prime in primes:
            heapq.heappush(minPQ, (prime, 0, prime))  # (value, index in output, prime factor)
        
        output = [1]  # The first super ugly number is always 1

        while len(output) < n:
            value, idx, prime = heapq.heappop(minPQ)
            
            # Avoid adding duplicate values to the output list
            if value != output[-1]:
                output.append(value)
            
            # Generate the next multiple of the current prime
            heapq.heappush(minPQ, (output[idx + 1] * prime, idx + 1, prime))
        
        return output[n - 1]

Java

class Solution {
    public int nthSuperUglyNumber(int n, int[] primes) {
        // Priority queue to get the smallest current number
        PriorityQueue minPQ = new PriorityQueue<>(Comparator.comparingLong(a -> a[0]));

        // Initialize the priority queue with the first prime multiples
        for (int prime : primes) {
            minPQ.offer(new long[] {prime, 0, prime}); // [value, index in output, prime factor]
        }

        // Array to store the ugly numbers
        long[] output = new long[n];
        output[0] = 1; // The first super ugly number is always 1

        // Populate the output array
        for (int i = 1; i < n; i++) {
            long[] current = minPQ.poll();
            long value = current[0];
            int idx = (int)current[1];
            long prime = current[2];

            // Avoid adding duplicate values to the output array
            if (value != output[i - 1]) {
                output[i] = value;
            } else {
                i--; // If duplicate, redo this index
            }

            // Generate the next multiple of the current prime
            minPQ.offer(new long[] {output[idx + 1] * prime, idx + 1, prime});
        }

        return (int)output[n - 1];
    }
}

JavaScript

/**
 * @param {number} n
 * @param {number[]} primes
 * @return {number}
 */
var nthSuperUglyNumber = function (n, primes) {
    const minPQ = new PriorityQueue({ compare: (a, b) => a[0] - b[0] });

    // Initialize the priority queue with the first prime multiples
    for (const prime of primes) {
        minPQ.enqueue([prime, 0, prime]); // [value, index in output, prime factor]
    }

    const output = [1]; // The first super ugly number is always 1

    while (output.length < n) {
        const [value, idx, prime] = minPQ.dequeue();
        
        // Avoid adding duplicate values to the output array
        if (value !== output[output.length - 1]) {
            output.push(value);
        }
        
        // Generate the next multiple of the current prime
        minPQ.enqueue([output[idx + 1] * prime, idx + 1, prime]);
    }

    return output[n - 1];
};