Question

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_ nth _super ugly number_.

The nth 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 <= 105
• 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)

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.

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

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