题解 | #最小的K个数#python三种方法实现sort、heap、partition

# -*- coding:utf-8 -*-
import heapq
#利用heapq实现一个大根堆
def heapify(x):
    for i in range(len(x)):
        x[i] = -x[i]
    heapq.heapify(x)

def heappush(heap, item):
    heapq.heappush(heap, -item)
def heappop(heap):
    return -heapq.heappop(heap)

class Solution:
    def GetLeastNumbers_Solution1(self, tinput, k):
        #直接快排 + 切片, O(nlog(n))
        if k == 0:
            return []

        return self.quicksort(tinput.copy(), 0, len(tinput)-1)[:k]

    def GetLeastNumbers_Solution2(self, tinput, k):
        #维护一个size为k的大根堆, 当堆满的时候，新的元素必须小于根节点才能插入，并弹出根节点
        #O(nlog(k))
        if k == 0:
            return []
        if k == len(tinput):
            return tinput

        heap_k = tinput[:k].copy()
        heapify(heap_k)
        for i in tinput[k:]:
            if i < -heap_k[0]:
                heappop(heap_k)
                heappush(heap_k, i)
        return [-e for e in heap_k]

    def GetLeastNumbers_Solution3(self, tinput, k):
        #利用快排分治的思想, O(n)
        if k == 0:
            return []
        if k == len(tinput):
            return tinput
        return self.get_least_k(tinput, 0, len(tinput)-1, k)[:k]

    def get_least_k(self, arr, low, high, k):
        #利用快排分治的思想
        i = low
        j = low
        pivot = high
        for j in range(low, high):
            if arr[j] < arr[pivot]:
                tmp = arr[j] 
                arr[j] = arr[i]
                arr[i] = tmp
                i += 1
        tmp = arr[i] 
        arr[i] = arr[pivot]
        arr[pivot] = tmp

        if i == k:
            return arr
        else:
            if i < k:
                arr = self.get_least_k(arr, i+1, high, k)
            else:
                arr = self.get_least_k(arr, low, i-1, k)
        return arr

    def quicksort(self, arr, low, high):
        if high - low < 2:
            return 
        i = low
        j = low
        pivot = high
        for j in range(low, high):
            if arr[j] < arr[pivot]:
                tmp = arr[j] 
                arr[j] = arr[i]
                arr[i] = tmp
                i += 1
        tmp = arr[i] 
        arr[i] = arr[pivot]
        arr[pivot] = tmp

        self.quicksort(arr, low, i-1)
        self.quicksort(arr, i+1, high)

        return arr

​