面试题 17.14. 最小K个数

题目

设计一个算法，找出数组中最小的k个数。以任意顺序返回这k个数均可。

示例：

输入： arr = [1,3,5,7,2,4,6,8], k = 4
输出： [1,2,3,4]

分析： 一种方法是创建一个长度为k的数组，保留最小的k个元素，从左往右遍历arr时，对当前元素与数组中元素做比较，替换掉大于当前元素的数组中的最大的那个元素，时间复杂度$O(n\ast k)$， 第二种方法是，先进行排序，然后返回前k个元素，时间复杂度$O(nlogn)$
方法1： 遍历，此方法超时

class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        n = len(arr)
        min_k = arr[:k]
        for i in range(k, n):
            max_num = arr[i]
            for j in range(k):
                if min_k[j] >= max_num:
                    min_k[j], max_num = max_num, min_k[j]
        return min_k

方法2： 快排， 同时可以进一步改进，我们知道快排会对数组进行二分为左右两部分，左部分小于右部分，其中通过下标来控制切分点，如果下标刚好等于k，则左部分刚好是最小的k个元素，如果大于，说明左部分数组包含最小的k个数，此时只需要递归左部分数组，反之小于，说明右部分存在最小的k个数，那么左右两部分均需要递归排序

class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        def quick_sort(i, j):
            if i >= j:
                return
            q = arr[j-1]
            c = i-1
            for m in range(i, j):
                if arr[m] <= q:
                    c += 1
                    arr[c], arr[m] = arr[m], arr[c]
            if c == k:
                return
            elif c < k:
                quick_sort(i, c)
                quick_sort(c+1, j)
            else:
                quick_sort(i, c)
        quick_sort(0, len(arr))
        return arr[:k]