LeetCode 面试题17.14.最小K个数

1.题目

原题链接

https://leetcode-cn.com/problems/smallest-k-lcci/

题目描述

设计一个算法，找出数组中最小的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))

2.算法

方法1 直接排序(超时)

思路与算法

直接将数组排序，返回前$K$个元素

复杂度

• 时间复杂度：$O(NlogN)$
• 空间复杂度：$O(1)$

方法2 快速选择

思路与算法

快速排序

代码

class Solution:
    def smallestK(self, arr, k):
        def partition(arr, lo, hi):
            pivot = arr[lo]
            while lo < hi:
                while lo < hi and arr[hi] >= pivot:
                    hi -= 1
                arr[lo] = arr[hi]
                while lo < hi and arr[lo] <= pivot:
                    lo += 1
                arr[hi] = arr[lo]
            arr[lo] = pivot
            return lo
        n = len(arr)
        if k >= n:
            return arr
        lo, hi = 0, n - 1
        while lo < hi:
            pos = partition(arr, lo, hi)
            if pos == k - 1:
                break
            elif pos < k - 1:
                lo = pos + 1
            else:
                hi = pos - 1
        return arr[:k]

复杂度

• 时间复杂度：$O(NlogK)$, 其中$K$可以认为是常数，故$O(N)$
• 空间复杂度：$O(1)$