313. Super Ugly Number – Leetcode Solutions

Example 1:

Input: n = 12, primes = [2,7,13,19]
Output: 32
Explanation: [1,2,4,7,8,13,14,16,19,26,28,32] is the sequence of the first 12 super ugly numbers given primes = [2,7,13,19].

Example 2:

Input: n = 1, primes = [2,3,5]
Output: 1
Explanation: 1 has no prime factors, therefore all of its prime factors are in the array primes = [2,3,5].

import heapq

class Solution:
    def nthSuperUglyNumber(self, n: int, primes: List[int]) -> int:
        # Initializing an array to store super ugly numbers
        ugly_numbers = [1]  # The first super ugly number is always 1
        
        # We will use a heap to keep track of the next smallest ugly number
        heap = []
        for prime in primes:
            # Push tuples of (prime, prime, 0) into the heap
            # Each tuple holds (value, prime, index) where:
            # - value is the current prime value
            # - prime is the prime number itself
            # - index is the index of the current ugly number being multiplied
            heapq.heappush(heap, (prime, prime, 0))
        
        # Loop until we get n ugly numbers
        while len(ugly_numbers) < n:
            # Get the smallest value from the heap
            next_ugly, prime, idx = heapq.heappop(heap)
            
            # If this ugly number is not a duplicate, add it to the result
            if next_ugly != ugly_numbers[-1]:
                ugly_numbers.append(next_ugly)
            
            # Push the next value of this prime multiplied by the next ugly number
            heapq.heappush(heap, (prime * ugly_numbers[idx + 1], prime, idx + 1))
        
        # The nth ugly number will be at index n-1 (since indexing starts at 0)
        return ugly_numbers[n - 1]

class Solution {
public:
    int nthSuperUglyNumber(int n, vector<int>& primes) {
        // Array to store the super ugly numbers
        vector<int> uglyNumbers(1, 1);  // The first super ugly number is 1
        
        // Min-heap (priority queue) to keep track of the next smallest ugly number
        priority_queue<tuple<long long, int, int>, vector<tuple<long long, int, int>>, greater<>> minHeap;
        
        // Initialize the heap with each prime number
        for (int prime : primes) {
            // Push tuple of (prime, prime, 0)
            // Each tuple holds (value, prime, index)
            // - value: current prime value (as long long to avoid overflow)
            // - prime: the prime number itself
            // - index: index of the current ugly number being multiplied
            minHeap.push(make_tuple((long long)prime, prime, 0));
        }
        
        // Continue generating ugly numbers until we have n numbers
        while (uglyNumbers.size() < n) {
            // Get the smallest value from the heap
            auto [nextUgly, prime, idx] = minHeap.top();
            minHeap.pop();
            
            // Only add to the list if it's not a duplicate
            if (nextUgly != uglyNumbers.back()) {
                uglyNumbers.push_back(nextUgly);
            }
            
            // Push the next multiple of this prime with the next ugly number
            minHeap.push(make_tuple(prime * (long long)uglyNumbers[idx + 1], prime, idx + 1));
        }
        
        // The nth super ugly number is the last element in the array
        return uglyNumbers[n - 1];
    }
};

var nthSuperUglyNumber = function(n, primes) {
    // Array to store the super ugly numbers
    const uglyNumbers = [1];  // The first super ugly number is always 1

    // Min-heap (priority queue) to keep track of the next smallest ugly number
    const minHeap = new MinHeap();

    // Initialize the heap with each prime number
    primes.forEach(prime => {
        // Each element in the heap is an object with:
        // - value: the current prime number
        // - prime: the prime itself
        // - index: index of the last ugly number multiplied with this prime
        minHeap.insert({ value: prime, prime: prime, index: 0 });
    });

    // Generate ugly numbers until we have n of them
    while (uglyNumbers.length < n) {
        let { value: nextUgly, prime, index } = minHeap.extractMin();

        // Only add the number if it's not a duplicate
        if (nextUgly !== uglyNumbers[uglyNumbers.length - 1]) {
            uglyNumbers.push(nextUgly);
        }

        // Insert the next multiple of the prime into the heap
        minHeap.insert({
            value: prime * uglyNumbers[index + 1],  // Next value is prime * the next ugly number
            prime: prime,
            index: index + 1
        });
    }

    // The nth super ugly number is the last number in the array
    return uglyNumbers[n - 1];
};

// Helper class to implement a Min-Heap (Priority Queue)
class MinHeap {
    constructor() {
        this.heap = [];
    }

    // Insert a new element into the heap
    insert(element) {
        this.heap.push(element);
        this.bubbleUp();
    }

    // Remove and return the smallest element from the heap
    extractMin() {
        if (this.heap.length === 1) {
            return this.heap.pop();
        }

        const min = this.heap[0];
        this.heap[0] = this.heap.pop();
        this.bubbleDown();
        return min;
    }

    // Move the last element up to its correct position
    bubbleUp() {
        let index = this.heap.length - 1;
        const element = this.heap[index];

        while (index > 0) {
            let parentIndex = Math.floor((index - 1) / 2);
            let parent = this.heap[parentIndex];

            if (element.value >= parent.value) break;

            this.heap[index] = parent;
            index = parentIndex;
        }

        this.heap[index] = element;
    }

    // Move the first element down to its correct position
    bubbleDown() {
        let index = 0;
        const length = this.heap.length;
        const element = this.heap[0];

        while (true) {
            let leftChildIndex = 2 * index + 1;
            let rightChildIndex = 2 * index + 2;
            let leftChild, rightChild;
            let swapIndex = null;

            if (leftChildIndex < length) {
                leftChild = this.heap[leftChildIndex];
                if (leftChild.value < element.value) {
                    swapIndex = leftChildIndex;
                }
            }

            if (rightChildIndex < length) {
                rightChild = this.heap[rightChildIndex];
                if (
                    (swapIndex === null && rightChild.value < element.value) ||
                    (swapIndex !== null && rightChild.value < leftChild.value)
                ) {
                    swapIndex = rightChildIndex;
                }
            }

            if (swapIndex === null) break;

            this.heap[index] = this.heap[swapIndex];
            index = swapIndex;
        }

        this.heap[index] = element;
    }
}

class Solution {
    public int nthSuperUglyNumber(int n, int[] primes) {
        // Array to store the super ugly numbers
        int[] uglyNumbers = new int[n];  // Array to hold the first n super ugly numbers
        uglyNumbers[0] = 1;  // The first super ugly number is always 1

        // PriorityQueue (min-heap) to keep track of the next smallest ugly number
        PriorityQueue<long[]> minHeap = new PriorityQueue<>((a, b) -> Long.compare(a[0], b[0]));

        // Initialize the heap with each prime number
        for (int prime : primes) {
            // Push tuple of (prime, prime, 0)
            // Each entry in the heap holds:
            // - value: current multiple of the prime
            // - prime: the prime number itself
            // - index: index of the last ugly number that the prime was multiplied with
            minHeap.offer(new long[]{prime, prime, 0});
        }

        // Generate super ugly numbers until we get n of them
        for (int i = 1; i < n; i++) {
            long[] current = minHeap.poll();  // Extract the smallest number from the heap
            long nextUgly = current[0];       // This is the next super ugly number
            int prime = (int) current[1];     // The prime associated with this value
            int index = (int) current[2];     // Index of the last ugly number used for this prime

            // Only add if it's not a duplicate
            if (nextUgly != uglyNumbers[i - 1]) {
                uglyNumbers[i] = (int) nextUgly;
            } else {
                i--;  // Avoid counting duplicates, decrement i so we don't skip adding a number
            }

            // Push the next multiple of this prime with the next ugly number
            long nextValue = prime * (long) uglyNumbers[index + 1];  // Calculate the next multiple
            minHeap.offer(new long[]{nextValue, prime, index + 1});
        }

        // The nth super ugly number is at index n-1
        return uglyNumbers[n - 1];
    }
}