面试题 17.09. 第 k 个数 :「优先队列」&「多路归并」

题目描述

这是 LeetCode 上的 面试题 17.09. 第 k 个数 ，难度为 困难。

Tag : 「优先队列（堆）」、「多路归并」、「双指针」

有些数的素因子只有 3，5，7，请设计一个算法找出第 k 个数。注意，不是必须有这些素因子，而是必须不包含其他的素因子。例如，前几个数按顺序应该是 1，3，5，7，9，15，21。

示例 1:

输入: k = 5

输出: 9
复制代码

基本分析

本题的基本思路与 264. 丑数 II : 从朴素优先队列到多路归并 完全一致。

优先队列（小根堆）

有了基本的分析思路，一个简单的解法是使用优先队列：

1. 起始先将最小数值 111 放入队列
2. 每次从队列取出最小值 xxx，然后将 xxx 所对应的数值 3x3x3x、5x5x5x 和 7x7x7x 进行入队
3. 对步骤 2 循环多次，第 kkk 次出队的值即是答案

为了防止同一数值多次进队，我们需要使用数据结构 SetSetSet 来记录入过队列的数值。

Java 代码：

class Solution {
    public int getKthMagicNumber(int k) {
        int[] nums = new int[]{3, 5, 7};
        PriorityQueue<Long> q = new PriorityQueue<>();
        Set<Long> set = new HashSet<>();
        q.add(1L); set.add(1L);
        while (!q.isEmpty()) {
            long t = q.poll();
            if (--k == 0) return (int) t;
            for (int x : nums) {
                if (!set.contains(x * t)) {
                    q.add(x * t); set.add(x * t);
                }
            }
        }
        return -1;
    }
}
复制代码

Python3 代码：

class Solution:
    def getKthMagicNumber(self, k: int) -> int:
        nums = [3, 5, 7]
        q, vis = [], set()
        q.append(1)
        vis.add(1)
        while q:
            t = heapq.heappop(q)
            k -= 1
            if k == 0:
                return t
            for x in nums:
                if t * x not in vis:
                    heapq.heappush(q, t * x)
                    vis.add(t * x)
        return -1
复制代码

• 时间复杂度：O(klog⁡k)O(k\log{k})O(klogk)
• 空间复杂度：O(k)O(k)O(k)

多路归并（多指针）

从解法一中不难发现，我们「往后产生的数值」都是基于「已有数值」而来（使用「已有数值」乘上 333、555、777）。

因此，如果我们最终的数值序列为 a1,a2,a3,...,ana1,a2,a3,...,ana1,a2,a3,...,an 的话，序列中的每一个数都必然能够被以下三个序列（中的至少一个）覆盖：

• 由数值 ×3\times 3×3 所得的有序序列：1×31 \times 31×3、2×32 \times 32×3、3×33 \times 33×3、4×34 \times 34×3、5×35 \times 35×3、6×36 \times 36×3、8×38 \times 38×3 ...
• 由数值 ×5\times 5×5 所得的有序序列：1×51 \times 51×5、2×52 \times 52×5、3×53 \times 53×5、4×54 \times 54×5、5×55 \times 55×5、6×56 \times 56×5、8×58 \times 58×5 ...
• 由数值 ×6\times 6×6 所得的有序序列：1×71 \times 71×7、2×72 \times 72×7、3×73 \times 73×7、4×74 \times 74×7、5×75 \times 75×7、6×76 \times 76×7、8×68 \times 68×6 ...

举个🌰，假设我们需要求得 [1,3,5,7,9,15,21][1, 3, 5, 7, 9, 15, 21][1,3,5,7,9,15,21] 数值序列 arrarrarr 的最后一位，那么该序列可以看作以下三个有序序列归并而来：

• 1×3,3×3,5×3,...,15×3,21×31 \times 3, 3 \times 3, 5 \times 3, ... , 15 \times 3, 21 \times 31×3,3×3,5×3,...,15×3,21×3 ，将 333 提出，即 arr×3arr \times 3arr×3
• 1×5,3×5,5×5,...,15×5,21×51 \times 5, 3 \times 5, 5 \times 5, ... , 15 \times 5, 21 \times 51×5,3×5,5×5,...,15×5,21×5 ，将 555 提出，即 arr×5arr \times 5arr×5
• 1×7,3×7,5×7,...,15×7,21×71 \times 7, 3 \times 7, 5 \times 7, ... , 15 \times 7, 21 \times 71×7,3×7,5×7,...,15×7,21×7 ，将 777 提出，即 arr×7arr \times 7arr×7

因此我们可以使用三个指针来指向目标序列 arrarrarr 的某个下标（下标 000 作为哨兵不使用，起始都为 111），使用 arr[下标]×系数arr[下标] \times 系数arr[下标]×系数（3、5 和 7） 代表当前使用到三个有序序列中的哪一位，同时使用 idxidxidx 表示当前生成到 arrarrarr 哪一位数值。

Java 代码：

class Solution {
    public int getKthMagicNumber(int k) {
        int[] ans = new int[k + 1];
        ans[1] = 1;
        for (int i3 = 1, i5 = 1, i7 = 1, idx = 2; idx <= k; idx++) {
            int a = ans[i3] * 3, b = ans[i5] * 5, c = ans[i7] * 7;
            int min = Math.min(a, Math.min(b, c));
            if (min == a) i3++;
            if (min == b) i5++;
            if (min == c) i7++;
            ans[idx] = min;
        }
        return ans[k];
    }
}
复制代码

TypeScript 代码：

function getKthMagicNumber(k: number): number {
    const ans = new Array<number>(k + 1).fill(1)
    for (let i3 = 1, i5 = 1, i7 = 1, idx = 2; idx <= k; idx++) {
        const a = ans[i3] * 3, b = ans[i5] * 5, c = ans[i7] * 7
        const min = Math.min(a, Math.min(b, c))
        if (a == min) i3++
        if (b == min) i5++
        if (c == min) i7++
        ans[idx] = min
    }
    return ans[k]
};
复制代码

Python 代码：

class Solution:
    def getKthMagicNumber(self, k: int) -> int:
        ans = [1] * (k + 1)
        i3, i5, i7 = 1, 1, 1
        for idx in range(2, k + 1):
            a, b, c = ans[i3] * 3, ans[i5] * 5, ans[i7] * 7
            cur = min([a, b, c])
            i3 = i3 + 1 if cur == a else i3
            i5 = i5 + 1 if cur == b else i5
            i7 = i7 + 1 if cur == c else i7
            ans[idx] = cur
        return ans[k]
复制代码

• 时间复杂度：O(k)O(k)O(k)
• 空间复杂度：O(k)O(k)O(k)

作者：宫水三叶的刷题日记
链接：https://juejin.cn/post/7148251901337796615
来源：稀土掘金
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。