TOP_K问题以及堆的应用

TOP_K问题

问题简介

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

三种解决方案

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

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

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

时间复杂度比较

# -*- 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())

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

# -*- 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

堆

构造堆以及堆排序

# -*- 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)

堆的应用 —— 最小的K个数

leetcode 题

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。
case1:
输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

根据题解中的思路leetcode题解
3. 大根堆方法
大根堆——父节点的值大于或等于子节点的值

主要看思路二
思路二: 将给定数组前 K 个数据进行大根堆 heap 的排序,然后从第 i = k+1 个元素开始跟大根堆的第一个元素,也就是当前 K个元素里面的最大值进行比较: 当 arr[i] < heap[0] 时,当堆外元素< 当前大根堆堆顶元素时，我们将交换这两个元素,重新进行大根堆的排序,如此进行 n-k次,我们的堆中就是我们最后要的结果.
这个思路的时间复杂度为O(NlogK),空间复杂度为O(K).此方法在海量数据时应用比较好,因为我们的内存有限不能将全部数据读入.

class Solution2:
    def getLeastNumbers(self, arr, k):
        # 使用大根堆的方法
        if k <= 0 or k > len(arr):
            return []
        heap = self.build_heap(arr[:k])
        for i in range(k,len(arr)):
            if arr[i] < heap[0]:
                # 交换
                heap[0] = arr[i]
                self.sink(heap,0)
        return heap

    # 大根堆
    def sink(self,array,k):
        n = len(array)
        left = 2*k + 1
        right = 2*k + 2
        if left >= n: return
        max_i = left
        if right < n and array[left] < array[right]:
            max_i = right
        # 根据前k个元素进行建堆
        if array[max_i] > array[k]:
            array[max_i],array[k] = array[k],array[max_i]
            self.sink(array,max_i)

    #建堆（堆化数组）
    def build_heap(self,list_):
        n = len(list_)
        for i in range(n//2,-1,-1):
            self.sink(list_,1)
        return list_
 if __name__ == "__main__":
    arr = [3,2,1]
    k = 2
    s2 = Solution2()
    print(s2.getLeastNumbers(arr, k)) #[2,1]