最小k个数

1,373 阅读3分钟

小知识,大挑战!本文正在参与“程序员必备小知识”创作活动

最小k个数

问题描述

给定一个长度为 n 的可能有重复值的数组,找出其中不去重的最小的 k 个数。例如数组元素是4,5,1,6,2,7,3,8这8个数字,则最小的4个数字是1,2,3,4(任意顺序皆可)。

要求:时间复杂度为O(nlogn),空间复杂度为O(n)。

示例:

输入:arr = [3,2,1], k = 2
输出:[1,2] 或者 [2,1]

分析问题

我们拿到这个问题,最先想到的就是对整个数组按照从小到大进行排序,然后取出前k个即可。

def smallestK(arr, k):
        arr.sort()
        return arr[:k]

arr=[4,5,1,6,2,7,3,8]
print(smallestK(arr,4))

复杂度分析:

  1. 时间复杂度:O(nlogn),其中n是数组的长度。该算法的时间复杂度就是排序的时间复杂度。
  2. 空间复杂度:O(logn)。排序的空间复杂度为O(logn)。

所以该算法是满足题目要求的,那我们还有更优的解决方案吗?我们来看一下。由于题目要求是求出前K个最小的数,并且不要求按照大小顺序输出,所以我们可以使用大顶堆来实现。

堆实现

开始时,我们将前k个数插入到大顶堆中,然后从第k+1个数开始遍历,如果当前遍历到的数比大顶堆的堆顶元素要小时,就把堆顶元素弹出,然后再把该数插入到堆中。最后将大顶堆中数返回即可。

Tips:在用python实现时,由于python中的堆为小顶堆,所以我们需要对数组中的所有数取相反数,才能使得采用小顶堆来维护前k小值。

import heapq
def smallestK(arr,k):
    if k == 0:
        return []
    n=len(arr)
    #前k个数取反
    data=[-x for x in arr[:k]]
    #建堆
    heapq.heapify(data)
    #print(data)
    for i in range(k, n):
        if -data[0] > arr[i]:
            heapq.heappop(data)
            heapq.heappush(data, -arr[i])
    ans = [-x for x in data]
    return ans

arr=[4,5,1,6,2,7,3,8]
print(smallestK(arr,4))

复杂度分析:

  1. 时间复杂度:O(nlogk),其中 n 是数组 arr 的长度。由于在大顶堆中维护k个元素,所以插入删除元素的时间复杂度为O(logK)。最坏的情况是,数组中的n个元素都会插入,所以该算法的时间复杂度是O(nlogK)。

  2. 空间复杂度:O(k),因为大顶堆中维持k个元素。

快排思想

其实,我们这里还可以借助快速排序的思想来求解这个问题。

因为题目不要求按照顺序输出,所以我们只需要把前k小的元素放到位置[0,k)即可,而快排正好可以做到这一点。快排每次都会将小于分区点的元素放在左边,大于分区点的元素放在右边。因此,我们可以通过判断分区点的下标q和k的关系来求解。

image.png

  1. q<k,表示分区点左侧元素不足k个,所以递归处理右边,让分区点下标右移。
  2. q>k,表示分区点左侧元素超过k个,所以需要递归处理左边,让分区点的下标左移。
  3. q=k,表示分区点左侧元素正好为k,输出分区点左侧元素即可。

下面我们来看一下代码实现。

import random
class Solution:
    def partition(self, nums, l, r):
        pivot = nums[r]
        i = l
        for j in range(l, r):
            if nums[j] < pivot:
                if i!=j:
                    temp = nums[i]
                    nums[i] = nums[j]
                    nums[j] = temp
                i = i + 1

        #将分区点放入相应位置
        temp=nums[i]
        nums[i]=nums[r]
        nums[r]=temp

        return i

    def select_partition(self, nums, l, r):
        #随机选取分区点的位置
        i = random.randint(l, r)
        #分区点放到最后一个元素
        temp = nums[r]
        nums[r] = nums[i]
        nums[i] = temp

        return self.partition(nums, l, r)

    def smallestK(self, arr, k):
        if k == 0:
            return list()
        n=len(arr)
        left=0
        right=n-1
        while left<right:
            index=self.partition(arr,left,right)
            if index==k-1:
                break
            elif index<k-1:
                left=index+1
            else:
                right=index-1
        return arr[:k]

s=Solution()
arr=[4,5,1,6,2,7,3,8]
count=4
print(s.smallestK(arr,count))