TOP_K问题以及堆排序

发表于 |

找出一个列表中前k大的数

TOP_K问题

问题简介

找出一个列表中前k大的数

三种解决方案

  • 排序后切片，时间复杂度O(nlogn+k)

  • 简单排序：冒泡排序为例，时间复杂度O(kn)

  • 小根堆，时间复杂度【建堆klogk + (n-k)logk = nlogk】

时间复杂度比较

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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
# -*- encoding: utf-8 -*-
"""
@File    : top_k.py
@Time    : 2020/1/15 12:38 下午
@Author  : zhengjiani
@Email   : 936089353@qq.com
@Software: PyCharm
"""
# 找前5个最大的,使用小根堆
# 使用内置heapq
import heapq

from leetcode.cal_time import cal_time

"""
1.取列表前k个元素建立一个小根堆，堆顶就是目前第K大的数
2.依次向后遍历原列表，对于列表中的元素，如果小于堆顶，则忽略该元素；如果大于堆顶，则将堆顶更换为该元素，并且对堆进行一次调整
3.遍历列表所有元素后，倒序弹出堆顶
"""
@cal_time
def func_1():
    li = [9,5,7,8,2,6,4,1,3]
    return heapq.nlargest(5,li)

@cal_time
def func_2():
    # 排序后切片
    li = [9, 5, 7, 8, 2, 6, 4, 1, 3]
    li.sort()
    return li[:3:-1]


@cal_time
def func_3():
    li = [9, 5, 7, 8, 2, 6, 4, 1, 3]
    _bubble_sort(li, 5)
    return li[:3:-1]


# 冒泡排序，只冒前5位
def _bubble_sort(li, k):
    for i in range(k):
        exchange = False
        for j in range(len(li) - i - 1):
            if li[j] > li[j + 1]:
                li[j], li[j + 1] = li[j + 1], li[j]
                exchange = True
        if not exchange:
            break

print(func_1())
print(func_2())
print(func_3())

# 控制台输出
/usr/local/bin/python3.7 /Users/zhengjiani/PycharmProjects/pyAlgorithm/leetcode/top_k.py
func_1 running time: 1.4066696166992188e-05 secs.
[9, 8, 7, 6, 5]
func_2 running time: 2.1457672119140625e-06 secs.
[9, 8, 7, 6, 5]
func_3 running time: 1.2159347534179688e-05 secs.
[9, 8, 7, 6, 5]

Process finished with exit code 0

== 计算时间的装饰器函数 ==

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# -*- encoding: utf-8 -*-
"""
@File    : cal_time.py
@Time    : 2020/1/3 5:51 下午
@Author  : zhengjiani
@Email   : 936089353@qq.com
@Software: PyCharm
"""
import time
def cal_time(func):
    def wrapper(*args, **kwargs):
        t1 = time.time()
        result = func(*args, **kwargs)
        t2 = time.time()
        print("%s running time: %s secs." % (func.__name__,t2-t1))
        return result
    return wrapper

堆

构造堆以及堆排序

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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
# -*- encoding: utf-8 -*-
"""
@File    : heap_sort.py
@Time    : 2020/1/14 2:22 下午
@Author  : zhengjiani
@Email   : 936089353@qq.com
@Software: PyCharm
"""
# 堆排序
# 大根堆
def sift(li,low,high):
    # li表示树，low表示树根，high表示树的最后一个节点的位置
    tmp = li[low]
    i = low
    j = 2 * i + 1 # 初始j指向空位的左孩子
    # i指向空位，j指向两个孩子
    while j <= high: # 循环退出的第二种情况：j>high,说明空位i是叶子节点
        if j+1 <= high and li[j] < li[j+1]: # 如果右孩子存在并且比左孩子大，指向右孩子
            j += 1
        if li[j] > tmp:
            li[i] = li[j]
            i = j
            j = 2 * i + 1
        else: # 循环退出的第一种情况：j位置的值都比tmp小，说明两个孩子都比tmp小
            break
    li[i] = tmp

# 小根堆
def sift_small(li,low,high):
    # li表示树，low表示树根，high表示树的最后一个节点的位置
    tmp = li[low]
    i = low
    j = 2 * i + 1 # 初始j指向空位的左孩子
    # i指向空位，j指向两个孩子
    while j <= high: # 循环退出的第二种情况：j>high,说明空位i是叶子节点
        if j+1 <= high and li[j] > li[j+1]: # 如果右孩子存在并且比左孩子大，指向右孩子
            j += 1
        if li[j] < tmp:
            li[i] = li[j]
            i = j
            j = 2 * i + 1
        else: # 循环退出的第一种情况：j位置的值都比tmp大，说明两个孩子都比tmp小
            break
    li[i] = tmp
    
def heap_sort_small(li):
    n = len(li)
    # 1.构造堆
    for low in range(n//2-1, -1, -1):
        sift_small(li, low, n-1)
    # 2.挨个出数
    for high in range(n-1, -1, -1):
        li[0], li[high] = li[high], li[0] # 退休棋子
        sift_small(li, 0, high-1)
        
# 构造堆,要求是完全二叉树，整个堆最后一个元素的位置当作high
def heap_sort(li):
    n = len(li)
    # 1.构造堆
    for low in range(n//2-1, -1, -1):
        sift(li, low, n-1)
    # 2.挨个出数
    for high in range(n-1, -1, -1):
        li[0], li[high] = li[high], li[0] # 退休棋子
        sift(li, 0, high-1)