堆heap的使用| 豆包MarsCode AI刷题

JoneLee
63 阅读5分钟

题目解析：找出出现频率前 k 高的元素

一、问题分析

给定一个整数数组 nums 和一个整数 k，我们需要找出数组中出现频率前 k 高的元素，并将这些元素按升序排列后返回。

输入说明：

  • nums 是一个正整数数组，长度最大为 10^5，数组中的元素范围是 [1, 10^4]。
  • k 是一个整数，表示我们需要返回频率前 k 高的元素。k 的值一定小于等于数组中不同元素的个数。

输出说明：

  • 返回一个包含频率前 k 高的元素的字符串，元素按升序排列，元素之间用逗号分隔。

二、解题思路

为了高效地解决这个问题，我们可以分为以下几步：

  1. 统计元素频率：我们首先需要统计数组中每个元素的出现频率。
  2. 选出前 k 高的频率：然后，我们需要选出频率最高的 k 个元素。由于 k 较小，我们可以利用堆（heap）来高效地维护这 k 个元素。
  3. 排序与格式化输出：最后，按照升序排列这些元素，并返回它们的字符串表示。

三、详细步骤及代码解析

1. 统计元素频率

我们首先通过 collections.Counter 来统计数组中每个元素的出现频率。Counter 是一个非常高效的工具，它可以直接生成一个字典，其中键是数组中的元素，值是该元素的出现次数。

from collections import Counter

freq_map = Counter(nums)

在这段代码中，freq_map 会是一个字典，键是数组中的元素，值是该元素的出现频率。例如，对于 nums = [1, 1, 1, 2, 2, 3]，freq_map 的结果会是：

Counter({1: 3, 2: 2, 3: 1})

2. 使用堆选出频率前 k 高的元素

接下来，我们需要使用堆来高效地选出频率前 k 高的元素。我们将使用 Python 的 heapq 库来实现最小堆。最小堆的特点是堆顶总是具有最小的元素，因此我们可以利用堆来维护一个大小为 k 的集合，使得堆顶是频率最小的元素。

我们将对频率和元素本身形成一个 (频率, 元素) 的元组，并将其推入堆中。如果堆的大小超过 k，则弹出堆顶的元素，从而保持堆中只保存频率前 k 高的元素。

import heapq

min_heap = []

# 将频率和元素的元组加入堆
for num, freq in freq_map.items():
    heapq.heappush(min_heap, (freq, num))
    
    # 如果堆的大小超过 k，则移除堆顶元素
    if len(min_heap) > k:
        heapq.heappop(min_heap)

在这段代码中，我们遍历 freq_map 中的每个元素和它的频率，使用 heapq.heappush 将其加入堆。然后，通过 heapq.heappop 删除堆顶元素，确保堆的大小始终保持为 k。

例如，如果 nums = [1, 1, 1, 2, 2, 3] 且 k = 2，堆的内容最终会是 (2, 2) 和 (3, 1)，表示频率前 2 高的元素是 1 和 2。

3. 从堆中提取元素并按升序排列

堆中存储的是元组 (频率, 元素)，我们最终关心的是元素本身。我们需要将堆中的元素按频率排序，并提取出元素部分。因为我们关心的是按升序排列的元素，因此需要对结果进行排序。

# 从堆中获取前 k 个高频元素
result = [num for freq, num in min_heap]

# 对结果进行升序排列
result.sort()

在这段代码中，我们通过列表推导式从堆中提取出元素部分，并使用 sort() 对结果进行升序排序。

4. 格式化输出

最后，我们将排序后的元素转换为字符串，并使用逗号分隔它们。我们可以使用 ','.join(map(str, result)) 来实现这个转换。

# 将结果转换为逗号分隔的字符串
return ','.join(map(str, result))

这段代码会将 result 列表中的元素转换为字符串，并用逗号连接它们。比如，如果 result = [1, 2]，那么返回的字符串就是 "1,2"。

四、完整代码

结合以上步骤，完整的解决方案代码如下：

import heapq
from collections import Counter

def solution(nums, k):
    # 统计每个元素的频率
    freq_map = Counter(nums)
    
    # 创建一个最小堆
    min_heap = []
    
    # 将频率和元素的元组加入堆
    for num, freq in freq_map.items():
        heapq.heappush(min_heap, (freq, num))
        
        # 如果堆的大小超过 k，则移除堆顶元素
        if len(min_heap) > k:
            heapq.heappop(min_heap)
    
    # 从堆中获取前 k 个高频元素
    result = [num for freq, num in min_heap]
    
    # 对结果进行升序排列
    result.sort()
    
    # 将结果转换为逗号分隔的字符串
    return ','.join(map(str, result))

五、MarsCode 帮助了我什么

MarsCode 提供了一个便捷的代码环境，让我快速验证了堆排序的实现。通过快速测试，我能够更好地理解堆的特性，以及如何使用 heapq 来优化此类问题的时间复杂度。MarsCode 的实时反馈和代码调试功能让我能够高效地优化解决方案。

六、总结

本题通过利用堆和哈希表的结合，高效地找出了频率前 k 高的元素，避免了排序整个数组，从而降低了时间复杂度。使用最小堆能够保证我们在遍历元素时，始终维护一个大小为 k 的集合，保证了算法的高效性。

该解法的时间复杂度主要由两个部分构成：

  1. 统计频率：O(n)，其中 n 是数组的长度。
  2. 使用堆维护频率前 k 高的元素：O(n log k)。

因此，总时间复杂度为 O(n log k)，这比直接排序 O(n log n) 更优。

avatar
文章
阅读
粉丝