最小K个数

题目描述

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

示例：

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

提示：

0 <= len(arr) <= 100000
0 <= k <= min(100000, len(arr))

分析与解答

想简单点，就是排序问题，然后输出前k个元素即可。我们要掌握的是排序算法，因为快速排序是面试最常考的排序算法，也是各个算法库基本都会内置的排序方法，因此，这里给出快速排序的实现。

快速排序思想：分治。
选定一个pivot，每一轮要保证pivot左侧都更小、右侧都更大。如何做到？指针对撞
[3, 1, 6, 2, 4, 5]选3为pivot，i=0，j=5；
j指针向左移动，5-4-3，j=3时arr[j]=2，这时让arr[i]=arr[j]，变为：[2, 1, 6, 2, 4, 5];
i指针向右移动，0-1-2，i=2时arr[i]=6，这时让arr[j]=arr[i]，变为：[2, 1, 6, 6, 4, 5];
j指针向左移动，3-2，j=2时i=j，这时让arr[i]=pivot，变为：[2, 1, 3, 6, 4, 5];
好了，现在对3的左侧和右侧分别进行同样的处理就可以了。

现在来看该思路具体实现：

class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        self.quickSort(arr, 0, len(arr) - 1)
        return arr[:k]

    # 快速排序，分治思想
    def quickSort(self, arr, left, right):
        if left < right:
            pivotIndex = self.partition(arr, left, right)
            self.quickSort(arr, left, pivotIndex - 1)
            self.quickSort(arr, pivotIndex + 1, right)

    # 每次调用partition()的结果就是pivot前面的元素都比pivot小、后面的元素都比pivot大
    def partition(self, arr, i, j):
        pivot = arr[i]

        while i < j:
            # j指针从右到左，直到找到小于pivot的元素，交换
            while i < j and arr[j] >= pivot:
                j -= 1
            arr[i] = arr[j]
            # i指针从左到右，直到找到大于pivot的元素，交换
            while i < j and arr[i] <= pivot:
                i += 1
            arr[j] = arr[i]

        # 最后i==j，这个位置是pivot，左侧都比pivot小、右侧都比pivot大
        arr[i] = pivot
        return i

技巧小结-快速排序

为了便于记忆，记住几个关键点，面试时回忆基本会想起来。
快速排序：pivot，分治（递归），指针对撞，交换。

拓展

当然可以直接这样：

arr.sort()
return arr[0:k]

或这样：

# heapq为最小堆
return heapq.nsmallest(k, arr)
# 当然，更好的方法是heapq始终维持k个元素然后全输出