313. Super Ugly Number

Description

A super ugly number is a positive integer whose prime factors are in the array primes.

Given an integer n and an array of integers primes, return the n^th super ugly number.

The n^th super ugly number is guaranteed to fit in a 32-bit signed integer.

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].

Constraints:

1 <= n <= 10⁵
1 <= primes.length <= 100
2 <= primes[i] <= 1000
primes[i] is guaranteed to be a prime number.
All the values of primes are unique and sorted in ascending order.

Solutions

Solution 1: Priority Queue (Min Heap)

Thinking

A super ugly number uses only the given primes. Trial-dividing every integer is too slow for large \(n\).

Start from \(1\), pop the heap minimum \(x\), and push \(x\times p\) when it does not overflow. If \(x\) is divisible by the current prime, skip later primes (Euler-sieve style). The \(n\)-th pop is the answer.

We use a priority queue (min heap) to maintain all possible super ugly numbers, initially putting \(1\) into the queue.

Each time we take the smallest super ugly number \(x\) from the queue, multiply \(x\) by each number in the array primes, and put the product into the queue. Repeat the above operation \(n\) times to get the \(n\)th super ugly number.

Since the problem guarantees that the \(n\)th super ugly number is within the range of a 32-bit signed integer, before we put the product into the queue, we can first check whether the product exceeds \(2^{31} - 1\). If it does, there is no need to put the product into the queue. In addition, the Euler sieve can be used for optimization.

The time complexity is \(O(n \times m \times \log (n \times m))\), and the space complexity is \(O(n \times m)\). Where \(m\) and \(n\) are the length of the array primes and the given integer \(n\) respectively.

Python3JavaC++Go

class Solution:
    def nthSuperUglyNumber(self, n: int, primes: List[int]) -> int:
        q = [1]
        x = 0
        mx = (1 << 31) - 1
        for _ in range(n):
            x = heappop(q)
            for k in primes:
                if x <= mx // k:
                    heappush(q, k * x)
                if x % k == 0:
                    break
        return x

class Solution {
    public int nthSuperUglyNumber(int n, int[] primes) {
        PriorityQueue<Integer> q = new PriorityQueue<>();
        q.offer(1);
        int x = 0;
        while (n-- > 0) {
            x = q.poll();
            while (!q.isEmpty() && q.peek() == x) {
                q.poll();
            }
            for (int k : primes) {
                if (k <= Integer.MAX_VALUE / x) {
                    q.offer(k * x);
                }
                if (x % k == 0) {
                    break;
                }
            }
        }
        return x;
    }
}

class Solution {
public:
    int nthSuperUglyNumber(int n, vector<int>& primes) {
        priority_queue<int, vector<int>, greater<int>> q;
        q.push(1);
        int x = 0;
        while (n--) {
            x = q.top();
            q.pop();
            for (int& k : primes) {
                if (x <= INT_MAX / k) {
                    q.push(k * x);
                }
                if (x % k == 0) {
                    break;
                }
            }
        }
        return x;
    }
};

func nthSuperUglyNumber(n int, primes []int) (x int) {
    q := hp{[]int{1}}
    for n > 0 {
        n--
        x = heap.Pop(&q).(int)
        for _, k := range primes {
            if x <= math.MaxInt32/k {
                heap.Push(&q, k*x)
            }
            if x%k == 0 {
                break
            }
        }
    }
    return
}

type hp struct{ sort.IntSlice }

func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any {
    a := h.IntSlice
    v := a[len(a)-1]
    h.IntSlice = a[:len(a)-1]
    return v
}

Solution 2: Dynamic Programming + Multi-pointer Heap

Thinking

Method 1 fans each ugly number across every prime and the heap grows large. Keep one pointer per prime: the heap top is the next candidate; after writing it, push that prime's next multiple. The heap stays \(O(m)\) and the time is \(O(n\log m)\).

Store the first \(n\) super ugly numbers in \(ugly[1..n]\), and keep a min-heap. Each heap entry belongs to one prime \(p\) and records the next candidate \(p \times ugly[\textit{index}]\).

Initialize \(ugly[1] = 1\) and push \((p, p, 2)\) for every prime. Repeatedly pop the heap minimum into \(ugly\), then advance that prime's pointer and push it back. Each prime keeps a single pointer, so we do not expand every generated ugly number against the whole prime list.

The time complexity is \(O(n \times \log m)\), and the space complexity is \(O(n + m)\), where \(m\) is the length of \(\textit{primes}\).

Go

type Ugly struct{ value, prime, index int }
type Queue []Ugly

func (u Queue) Len() int           { return len(u) }
func (u Queue) Swap(i, j int)      { u[i], u[j] = u[j], u[i] }
func (u Queue) Less(i, j int) bool { return u[i].value < u[j].value }
func (u *Queue) Push(v any)        { *u = append(*u, v.(Ugly)) }
func (u *Queue) Pop() any {
    old, x := *u, (*u)[len(*u)-1]
    *u = old[:len(old)-1]
    return x
}

func nthSuperUglyNumber(n int, primes []int) int {
    ugly, pq, p := make([]int, n+1), &Queue{}, 2
    ugly[1] = 1
    heap.Init(pq)
    for _, v := range primes {
        heap.Push(pq, Ugly{value: v, prime: v, index: 2})
    }
    for p <= n {
        top := heap.Pop(pq).(Ugly)
        if ugly[p-1] != top.value {
            ugly[p], p = top.value, p+1
        }
        top.value, top.index = ugly[top.index]*top.prime, top.index+1
        heap.Push(pq, top)
    }
    return ugly[n]
}

313. Super Ugly Number

Description

Solutions

Solution 1: Priority Queue (Min Heap)

Solution 2: Dynamic Programming + Multi-pointer Heap

Comments