题目描述

编写一段程序来查找第 n 个超级丑数。

超级丑数是指其所有质因数都是长度为 k 的质数列表 primes 中的正整数。

样例

输入：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]。

限制

• 1 是任何给定 primes 的超级丑数。
• 给定 primes 中的数字以升序排列。
• 0 < k <= 100, 0 < n <= 10^6, 0 < primes[i] < 1000。
• 第 n 个超级丑数确保在 32 位有符整数范围内。

算法1

(枚举) $O(nk)$

1. 定义一个长度为 n 的结果数组 res，第一个位置放 1，我们按顺序产生前 n 个数字。
2. 再初始化一个长度为 k 的下标数组，代表每个质数下一次与 res 数组元素作乘积的下标。初始时，每个下标都指向 res 数组的 0 位置，表示下一次都要与 res[0] 作乘积。
3. 枚举所有的质数，令所有质数与其对应位置作乘积，找到最小的乘积，就是下一个产生的数字。
4. 将所有能产生这个数字的质数的下标都向后移动一位。

时间复杂度

• 产生一个数字的时间复杂度为 $O(k)$，故总时间复杂度为 $O(nk)$。

空间复杂度

• 需要 $O(n + k)$ 的额外空间存储 res 数组和下标数组。

C++ 代码

class Solution {
public:
    int nthSuperUglyNumber(int n, vector<int>& primes) {
        int k = primes.size();
        vector<int> res(n), idx(k, 0);

        res[0] = 1;
        for (int i = 1; i < n; i++) {
            int mi = INT_MAX;
            for (int j = 0; j < k; j++) {
                if (1ll * res[idx[j]] * primes[j] > INT_MAX)
                    continue;

                mi = min(mi, res[idx[j]] * primes[j]);
            }

            res[i] = mi;
            for (int j = 0; j < k; j++) {
                if (1ll * res[idx[j]] * primes[j] > INT_MAX)
                    continue;

                if (mi == res[idx[j]] * primes[j])
                    idx[j]++;
            }
        }

        return res[n - 1];
    }
};

算法2

(优先队列) $O(n \log k)$

1. 将算法 1 替换为优先队列寻找最小值，但有可能每次有多个下标出队（产生的数字重复），造成时间复杂度上升。
2. 但重复的次数并不多，算法的常数很小，在 k 很大的时候表现比较优秀。

时间复杂度

• 初始化优先队列的时间复杂度
• 产生一个数字的时间复杂度为 $O(\log k)$，故总时间复杂度为 $O(n \log k)$。
• 但对于部分数据，即每次有多个质数产生了相同的新数字时，会造成时间复杂度上升。

空间复杂度

• 需要 $O(n + k)$ 的额外空间存储 res 数组，下标数组和优先队列。

C++ 代码

class Solution {
public:
    int nthSuperUglyNumber(int n, vector<int>& primes) {
        #define PI pair<int, int>

        int k = primes.size();
        vector<int> res(n), idx(k, 0);
        priority_queue<PI, vector<PI>, greater<PI>> heap;

        res[0] = 1;
        for (int i = 0; i < k; i++)
            heap.push(make_pair(primes[i], i));

        for (int i = 1; i < n; i++) {
            int mi = heap.top().first;
            res[i] = mi;

            while (heap.top().first == mi) {
                int p = heap.top().second;
                heap.pop();
                idx[p]++;

                if (1ll * res[idx[p]] * primes[p] <= INT_MAX)
                    heap.push(make_pair(res[idx[p]] * primes[p], p));
            }
        }

        return res[n - 1];
    }
};