Jianbo Dai's

jz040.最小的k个数

 
  • Jul 06, 2020

题目链接

这道题最简单的思路莫过于把输入的n个整数排序，排序之后位于最前面的k个数就是最小的k个数，这种思路的时间复杂度为O（NlogN）。

但是上面的算法并不是最快的，还有一种思想，可以从jz039中得到启发，我们也同样可以基于Partition函数来解决问题。如果基于数组的第k个数字来调整，则使得比第k个数字小的所有数字都位于数组的左边，比第k个数字大的所有数字都在数组的右边。这样调整之后，位于数组中左边的k个数字就是最小的k个数字（这k个数字不一定是排序的）。代码如2所示，随机选择pivot可以让平均情况更快，否则容易被某些特例拖长时间。这个算法的时间复杂度是O（N），空间复杂度由于是在原地排序所以是O（1）。

第三个方法是用堆来实现，见3，之后来更新。

import random

class Solution:
    def getLeastNumbers1(self, arr: List[int], k: int) -> List[int]:
        arr = sorted(arr)
        return arr[:k]

    def getLeastNumbers2(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return []
        t = k - 1
        l, r = 0, len(arr)-1
        while True:
            index = self.qsort(arr, l, r)
            if index == t:
                return arr[:index+1]
            elif index > t:
                r = index - 1
            else:
                l = index + 1

    def qsort(self, lst, left, right):
        random_index = random.randint(left, right) 
        lst[left], lst[random_index] = lst[random_index], lst[left]
        pivot = lst[left]
        i, j = left, right
        while i < j:
            while i < j and lst[j] >= pivot:
                j -= 1
            lst[i] = lst[j]
            while i < j and lst[i] <= pivot:
                i += 1
            lst[j] = lst[i]
        lst[i] = pivot
        return i

    def getLeastNumbers3(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return list()
        hp = [-x for x in arr[:k]]
        heapq.heapify(hp)
        for i in range(k, len(arr)):
            if -hp[0] > arr[i]:
                heapq.heappop(hp)
                heapq.heappush(hp, -arr[i])
        ans = [-x for x in hp]
        return ans
Search