第 k 个数

标签: 哈希表 数学 动态规划 堆(优先队列)

难度: Medium

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

示例 1:

输入: k = 5

输出: 9

Submission

运行时间: 24 ms

内存: 16.6 MB

class Solution:
    def getKthMagicNumber(self, k: int) -> int:
        heap = [1]
        seen = {1}
        factors = [3, 5, 7]

        for i in range(k-1):
            cur = heapq.heappop(heap)
            for x in factors:
                if (next := x * cur) not in seen:
                    seen.add(next)
                    heapq.heappush(heap, next)
        return heapq.heappop(heap)                    

Explain

该题解使用了最小堆(优先队列)以及哈希集合来求解。首先,1被视为第一个满足条件的数,因此初始化堆和集合都包含这个数。对于每次从堆中弹出的最小数(当前的最小数),将它与3、5、7分别相乘,得到的新数如果之前未出现过,就加入到堆和集合中。这个过程重复k次,最后从堆中弹出的数即为第k个数。

时间复杂度: O(k log k)

空间复杂度: O(k)

class Solution:
    def getKthMagicNumber(self, k: int) -> int:
        heap = [1]  # 初始化堆,开始时只有元素1
        seen = {1}  # 初始化集合,记录已经看到的元素,避免重复处理
        factors = [3, 5, 7]  # 只包含3,5,7的因子列表

        for i in range(k-1):  # 需要弹出k-1个元素后,堆顶元素即第k个数
            cur = heapq.heappop(heap)  # 从堆中弹出最小元素
            for x in factors:
                next = x * cur  # 生成新的数
                if next not in seen:  # 如果这个数是新的
                    seen.add(next)  # 加入集合
                    heapq.heappush(heap, next)  # 加入堆中
        return heapq.heappop(heap)  # 返回堆顶元素,即第k个数

Explore

使用最小堆(优先队列)的原因在于该数据结构提供了高效地访问和删除最小元素的能力,这对于本题是必需的,因为我们需要按顺序找到第k个最小的合格数。最小堆能够保证每次从堆顶弹出的总是当前堆中最小的数,这是数组或链表无法高效完成的。如果使用数组,每次寻找最小数需要O(n)时间,使用链表虽然可以在O(1)时间内添加元素,但删除最小元素平均需要O(n)时间。相比之下,最小堆可以在O(log n)时间内完成插入和删除最小元素的操作,更适合解决本题。

使用哈希集合来检查元素是否已存在是一个效率很高的方法。哈希集合在平均情况下提供O(1)时间复杂度的查找、插入和删除操作。这意味着即使对于大量的数据,哈希集合都能快速响应是否有重复元素的查询。在本题中,考虑到需要频繁地检查和插入新生成的数,使用哈希集合可以显著减少时间开销,相比于其他如列表等需要O(n)时间复杂度的数据结构,哈希集合更有效。因此,哈希集合通常不会成为性能瓶颈。

堆的最大大小可能达到3k是因为每次处理时,我们都可能向堆中添加最多三个新的数(当前最小数乘以3、5和7)。即使每次弹出堆顶元素,依然可能在每次迭代中净增加两个元素(如果所有生成的新数都是唯一的)。在进行k次迭代的过程中,最坏的情况是每次迭代都将三个新元素添加到堆中,因此堆的大小可能在某个时刻达到3k。虽然实际情况下由于重复的数不会被重新添加,堆的大小通常会小于这个理论最大值,但为确保算法的正确性和安全性,我们考虑最坏情况下的堆大小。