03-排序 topk

一、TopK 问题定义

topK问题比较经典，在面试算法题中考察也比较普遍，就是在一群无序的数中，找到最小的K个数和最大的K个数，或者第K个最大的数和第K个最小的数

本篇以寻找最小的K个数为例

二、解法

1.快速排序

面对topK问题，最能直接了当想到的做法便是排序，而后根据顺序用下标访问该元素即可。

  • 平均时间复杂度：O(nlogn)
  • 空间复杂度：O(1)

2.冒牌排序优化版

既然只要前 K 的数据，那我们可以使用冒泡排序只循环前 k 次，选出前 k 个值 

1
2
3
4
5
6
def bubble_sort_topk(li, k):
    for i in range(k):
        for j in range(len(li) - 1, i, -1):
            if li[j] > li[j - 1]:
                li[j], li[j - 1] = li[j - 1], li[j]
    return li[0: k]

3.快速排序优化版

其实快排中我们将数组分为了两部分，一部分是大的值，一部分是小的值，我们要找的第K个数，没有必要对整个数组进行寻找，可以对其中一侧进行寻找，从而达到优化的效果

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def partition(li, left, right):
    tmp = li[left]
    while left < right:
        while right > left and li[right] >= tmp:
            right -= 1
        li[left] = li[right]

        while left < right and li[left] <= tmp:
            left += 1
        li[right] = li[left]

    li[left] = tmp
    return left


def quick_sort_topyK(li, left, right, k):
    li_num = len(li)
    if left < right:
        mid = partition(li, left, right)
        if mid > k:
            quick_sort_topyK(li, left, mid)
            # quick_sort_topyK(li, mid + 1, right)
        else:
            # quick_sort_topyK(li, left, mid)
            quick_sort_topyK(li, mid + 1, right)


def quick_sort_topyK_1(array, k):
    start, end = 0, len(array) - 1
    index = partition(array, start, end)
    while index != k:
        if index < k:
            index = partition(array, index+1, end)
        else:
            index = partition(array, start, index)
    print(array[0: k])
  • 平均时间复杂度：O(n)
  • 空间复杂度：O(1)

4.堆排序

其实寻找到第 K 个数，同样可以通过维护一个堆来解决，

我们知道，堆顶要么存放最大数（大根堆），要么存放最小数（小根堆），

  • 当要找前 K 小的数值时

    • 取列表 K 个元素建立一个大根堆后堆顶就是目前第 K 大的数。
    • 依次向后遍历原列表，对于列表中的元素，如果大于堆顶，则忽略该元素 如果小于堆顶，则将堆顶更换为该元素，并且对堆进行一次向下调整
    • 最后遍历列表所有元素后，倒序弹出堆顶

  • 当要找前 K 大的数值时

    • 取列表 K 个元素建立一个小根堆后堆顶就是目前第 K 大的数。
    • 依次向后遍历原列表，对于列表中的元素，如果小于堆顶，则忽略该元素 如果大于堆顶，则将堆顶更换为该元素，并且对堆进行一次向下调整
    • 最后遍历列表所有元素后，倒序弹出堆顶

寻找前 K 小代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def sift(li, low, high):
    tmp = li[low]
    i = low
    j = i * 2 + 1
    while j <= high:
        if j + 1 <= high and li[j + 1] > li[j]:
            j = j + 1
        if li[j] > tmp:
            li[i] = li[j]
            i = j
            j = i * 2 + 1
        else:
            break

    li[i] = tmp


def heap_sort_topk(li, k):
    li_len = len(li)

    for i in range((k - 2) // 2, -1, -1):
        sift(li, i, k-1)
    for i in range(k, li_len-1):
        if li[i] < li[0]:
            li[i], li[0] = li[0], li[i]
            sift(li, 0, k-1)
    print(li[0: k])


def heap_sort_topk_order(li, k):
    li_len = len(li)

    for i in range((k -2) // 2, -1, -1):
        sift(li, i, k-1)

    for i in range(k, li_len):
        if li[i] < li[0]:
            li[i], li[0] = li[0], li[i]
            sift(li, 0, k-1)

    for i in range(k - 1, -1, -1):
        li[0], li[i] = li[i], li[0]
        sift(li, 0, i - 1)
    print(li)

寻找前 k 大代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
# 堆的向下调整函数
def sfit(li, low, height):
    i = low         # i 指向堆顶的位置
    j = i * 2 + 1   # j 目前指向i的右孩子的位置
    tmp = li[low]   # tmp 等于当前堆顶元素
    while j <= height:
        if j + 1 <= height and li[j + 1] < li[j]:
            j = j + 1
        if li[j] < tmp:
            li[i] = li[j]
            i = j   # i 指向第二层（左边或右边）的子树的堆顶
            j = i * 2 + 1
        else:
            break   # 如果右孩子或者左孩子不大于tmp则循环退出
    li[i] = tmp
 

def heap_sort_topk(li, k):
    heap = li[0: k]
    # 1、建立小根堆
    for i in range((k - 2) // 2, -1, -1):
        sfit(heap, i, k - 1)
    # 2、拿堆顶与后面的元素对比
    for i in range(k, len(li) - 1):
        if li[i] > heap[0]:
            heap[0] = li[i]
            sfit(heap, 0, k - 1)
    # 3.挨个出数，使得列表倒序
    for i in range(k - 1, -1, -1):
        heap[0], heap[i] = heap[i], heap[0]
        sfit(heap, 0, i - 1)
    return heap
  • 平均时间复杂度：O(nlog(n))
  • 空间复杂度：O(n)

对比

在海量数据中，如何快速找到TopK个数据？首先以上三个算法都能做到，但为什么堆排序是最优。

在快排与冒泡中, 对于数组 nums，空间复杂度为 O(N), 也就是说一次性都要读入内存中， 先将 nums 全部排好，然后取出最大的 topk 个数字。当数据量过大时，一台机器是没有办法读取的, 而且时间复杂度上是 O(NlogN)。

而堆排序，它只要维护一个大小为 K 的大顶堆, 依次将 nums 的数入堆，当堆大小为 K 满了,则下一个数 number 入堆时，只需要将堆顶元素与下一个数 number 比较。number 大于堆顶元素，则将当前的堆顶元素抛弃，并将该元素插入堆中，重新进行堆排序，永远保证堆顶元素是 K 个数里最小的。

因此，这个堆排序的时间复杂度是 O(Nlog(K))， 而当 nums 特别大时，远大于 K, 那么时间复杂度就是 O(N), 堆排的复杂度可忽略，而且这样内存中只用维护大小为 K 的堆，然后每次读取一个 number，是不会存在超内存情况的，所以实践中 TopK 问题用堆排序是最优的。

算法
#算法
03-排序 topk
https://flepeng.github.io/030-算法-03-排序-topk/
作者
Lepeng
发布于
2020年8月8日
许可协议