解法一:基于Partidion函数的平均时间复杂度为O(n)的方法。最坏O(kn),此时数组已经有序,需要执行Partidion函数k次。最好情况O(n),执行一次Partidion函数即得出结果。
Python3代码如下(最小的k个数):
def Partidion(nums,left,right):
p = nums[left]
while left < right:
while left < right and nums[right] >= p:
right -= 1
nums[left] = nums[right]
while left < right and nums[left] <= p:
left += 1
nums[right] = nums[left]
nums[left] = p
return left
def s(nums,k,r):
left,right = 0,len(nums)-1
p = Partidion(nums,left,right)
while p+1 != k:
if p+1 < k:
left = p+1
p = Partidion(nums,left,right)
else:
right = p-1
p = Partidion(nums,left,right)
for i in range(k):
r.append(nums[i])
return r
print(s([4,2,3,0,5,-1,6,-2,1],3,[]))
解法二:基于堆的平均时间复杂度为O(nlogn)。
Python3代码如下(最大的k个数):
import heapq
def f(nums,k):
h = []
for i in range(k):
h.append(nums[i])
heapq.heapify(h)
for j in range(k,len(nums)):
if nums[j] > h[0]:
heapq.heappop(h)
heapq.heappush(h,nums[j])
return h
print(f([0,2,3,0,5,6,0],3))