40-最小的k个数

题目

思路

思路一：排序

排序，快速排序，然后取出前k个。最大时间复杂度O(nlogn)

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        arr.sort()
        return arr[:k]

思路二：计数排序

维护一个数组，该数组共max(arrd) + 1个元素，初始化每个元素都为0，索引代表值，元素表示相应值出现的次数。

遍历一遍原始数组arr，得到每个数出现的次数。
在顺序遍历一遍计数数组，取出前k个数。
最大时间复杂度O(n)。因为每次遍历都是O(n)。

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        # 计数排序
        countNums = [0] * (max(arr) + 1)
        for num in arr:
            countNums[num] += 1
        ans = []
        for num, cnt in enumerate(countNums):
            if k == 0:
                break
            if k - cnt >= 0:
                ans += [num] * cnt
                k -= cnt
            else:
                ans += [num] * k
                k = 0
        return ans

思路三：基于快速排序的思路

快速排序中，每次选择一个pivot，然后partition将原始数组分成两部分，左半部分小于等于pivot，右半部分大于等于pivot。
如果该pivot所在的索引正好是k，那么前k个就是最小的k个数。否则，左半部分去找或者右半部分去找。

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        # target index is k
        # quick sort
        size = len(arr)
        if size == 0 or size <= k:
            return arr
        left = 0
        right = size - 1
        while True:
            # print(arr)
            ind = self._partition(arr, left, right)
            if ind == k:
                return arr[:k]
            elif ind > k:
                right = ind - 1
            else:
                left = ind + 1
    def _partition(self, arr, left, right):
        pivot = arr[left]
        j = left
        for i in range(left+1, right+1):
            if pivot > arr[i]:
                j += 1
                arr[i], arr[j] = arr[j], arr[i]
        arr[left], arr[j] = arr[j], arr[left]
        return j

思路四：基于堆或者红黑树的思路