面试题40. 最小的k个数
输入整数数组 arr ,找出其中最小的 k 个数。
例如,输入4、5、1、6、2、7、3、8这8个数字,则最小的4个数字是1、2、3、4。
示例 1:
输入:arr = [3,2,1], k = 2
输出:[1,2] 或者 [2,1]
示例 2:
输入:arr = [0,1,2,1], k = 1
输出:[0]
5分钟想到的代码是这样的
排序,前k个数打包,送走。于是一行代码就出来了
class Solution:
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
return(sorted(arr)[0:k])
然后进入改进模式(一个小时)
首先想到的是排序是大排序,那么如果一个一个进,每次进去都弹出最大值呢,这样不动原数组就能解决了,于是,,,,,于是就写了一个更复杂的代码出来(同行看了想打人)
class Solution:
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
if(k==0):
return []
mid=0
c=arr[0:k]
mid=max(c)
front=c.index(mid)
for i in range(k,len(arr)):
if(arr[i]>=mid):
continue
else:
del c[front]
c.append(arr[i])
mid=max(c)
front=c.index(mid)
return c
你为什么要折磨CPU.jpg
不过好在想起来了python有堆这种堆顶一直为最小值的结构,所以如果我们取相反数把列表做成堆,那么就不用特地设一个指针来表明最大值了,代码如下
class Solution:
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
if k == 0:
return list()
hp = [-i for i in arr[:k]]
heapq.heapify(hp)
for i in range(k, len(arr)):
if -hp[0] > arr[i]:
heapq.heappop(hp)
heapq.heappush(hp, -arr[i])
ans = [-i for i in hp]
return ans
果然是直观上越简单的题改进的越多(点头)