Question

313. 超级丑数

English Version

题目描述

超级丑数 是一个正整数，并满足其所有质因数都出现在质数数组 primes 中。

给你一个整数 n 和一个整数数组 primes ，返回第 n 个 超级丑数 。

题目数据保证第 n 个 超级丑数 在 32-bit 带符号整数范围内。

 

示例 1：

输入：n = 12, primes = [2,7,13,19]
输出：32 
解释：给定长度为 4 的质数数组 primes = [2,7,13,19]，前 12 个超级丑数序列为：[1,2,4,7,8,13,14,16,19,26,28,32] 。

示例 2：

输入：n = 1, primes = [2,3,5]
输出：1
解释：1 不含质因数，因此它的所有质因数都在质数数组 primes = [2,3,5] 中。

 

提示：

• 1 <= n <= 105
• 1 <= primes.length <= 100
• 2 <= primes[i] <= 1000
• 题目数据 保证 primes[i] 是一个质数
• primes 中的所有值都 互不相同 ，且按 递增顺序 排列

解法

方法一：优先队列（小根堆）

我们用一个优先队列（小根堆）维护所有可能的超级丑数，初始时将 $1$ 放入队列中。

每次从队列中取出最小的超级丑数 $x$，将 $x$ 乘以数组 primes 中的每个数，将乘积放入队列中，然后重复上述操作 $n$ 次即可得到第 $n$ 个超级丑数。

由于题目保证第 $n$ 个超级丑数在 $32$ 位带符号整数范围内，因此，我们将乘积放入队列之前，可以先判断乘积是否超过 $2^{31} - 1$，如果超过，则不需要将乘积放入队列中。另外，可以使用欧拉筛优化。

时间复杂度 $O(n \times m \times \log (n \times m))$，空间复杂度 $O(n \times m)$。其中 $m$ 和 $n$ 分别为数组 primes 的长度和给定的整数 $n$。

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
}

方法二

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