小知识,大挑战!本文正在参与“程序员必备小知识”创作活动
最小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))
复杂度分析:
- 时间复杂度:O(nlogn),其中n是数组的长度。该算法的时间复杂度就是排序的时间复杂度。
- 空间复杂度: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))
复杂度分析:
-
时间复杂度:O(nlogk),其中 n 是数组 arr 的长度。由于在大顶堆中维护k个元素,所以插入删除元素的时间复杂度为O(logK)。最坏的情况是,数组中的n个元素都会插入,所以该算法的时间复杂度是O(nlogK)。
-
空间复杂度:O(k),因为大顶堆中维持k个元素。
快排思想
其实,我们这里还可以借助快速排序的思想来求解这个问题。
因为题目不要求按照顺序输出,所以我们只需要把前k小的元素放到位置[0,k)即可,而快排正好可以做到这一点。快排每次都会将小于分区点的元素放在左边,大于分区点的元素放在右边。因此,我们可以通过判断分区点的下标q和k的关系来求解。
- q<k,表示分区点左侧元素不足k个,所以递归处理右边,让分区点下标右移。
- q>k,表示分区点左侧元素超过k个,所以需要递归处理左边,让分区点的下标左移。
- 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))