数据结构与算法:top K 问题

文章目录

    • 1. 找到数组中最小的 k 个数
      • 1.1 快排,O(nlogn),O(logn)
      • 1.2 大根堆,O(nlogk),O(k)

top K 问题是面试中常考的问题,往往可以用排序(排序)和堆(大/小根堆)来解决。当然,如果你只会快排的话,就可以早点面试结束回家吃饭了。

结论:

top K 问题常用堆来解决

1. 找到数组中最小的 k 个数

题目:

设计一个算法,找出数组中最小的k个数。以任意顺序返回这k个数均可。

示例:

输入: arr = [1,3,5,7,2,4,6,8], k = 4
输出: [1,2,3,4]

1.1 快排,O(nlogn),O(logn)

最容易想到的方法是快速排序,其时间复杂度为 O(nlogn),空间复杂度为 O(logn)

code:

class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        arr.sort()
        return arr[:k]

1.2 大根堆,O(nlogk),O(k)

我们用一个大根堆实时维护数组的前 k 小值。首先将前 k 个数插入大根堆中,随后从第 k+1 个数开始遍历,如果当前遍历到的数比大根堆的堆顶的数要小,就把堆顶的数弹出,再插入当前遍历到的数。最后将大根堆里的数存入数组返回即可。

由于 Python 中的 heapq库创建的堆为小根堆,因此我们要对数组中所有的数取其相反数变成大根堆。

code:

class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        if not arr or k==0:
            return []
        heap=[-x for x in arr[:k]]
        heapq.heapify(heap)
        for x in arr[k:]:
            if -heap[0]>x:
                heapq.heappop(heap)
                heapq.heappush(heap,-x)
        return [-x for x in heap]

复杂度分析:

  • 时间复杂度:O(nlogk),其中 n 是数组 arr 的长度。由于大根堆实时维护前 k 小值,所以插入删除都是 O(logk) 的时间复杂度,最坏情况下数组里 n 个数都会插入,所以一共需要 O(nlogk) 的时间复杂度。
  • 空间复杂度:O(k),因为大根堆里最多有 k 个数

相关题目持续更新中 …

你可能感兴趣的:(数据结构和算法,LeetCode,数据结构,算法,排序算法)