第 40 题：最小的 K 个数
传送门：最小的k个数，牛客网 online judge 地址。
输入
个整数，找出其中最小的
个数。
注意：
- 数据保证 k 一定小于等于输入数组的长度;
- 输出数组内元素请按从小到大顺序排序;
样例：
输入：
[1,2,3,4,5,6,7,8],k=4输出：
[1,2,3,4]
分析：最简单的思路就是排个序，然后取出前 个元素，不过时间复杂度是
，
为数组的长度。
Python 代码：
class Solution(object):
def getLeastNumbers_Solution(self, input, k):
size = len(input)
if size == 0:
return []
if k == size:
return sorted(input)
return sorted(input)[:k]
其实可以用 的时间复杂度找到最小的
个元素。有两种思路：
1、借助快速排序把数组一分为二的 partition 操作；
2、借助最大堆（需要把数组做一个转换，都变成相反数，最小的 个数，就是最大堆里最大的
个数）。
Python 代码1：partition，注意，这种方式，大于等于 pivot 的元素都被分在了右边
class Solution(object):
def getLeastNumbers_Solution(self, input, k):
"""
:type input: list[int]
:type k: int
:rtype: list[int]
"""
size = len(input)
if size == 0:
return []
if k == size:
return sorted(input)
l = 0
r = size - 1
while l <= r:
p = self.__partition(input, l, r)
if p == k - 1:
return sorted(input[:p + 1])
elif p > k - 1:
# 此时 k-1 p
r = p - 1
else:
# 此时 p k-1
l = p + 1
def __partition(self, input, left, right):
# 只有一个数的时候，就没有必要 partition 了
# 直接返回这个数的索引
if left == right:
return left
pivot = input[left]
j = left
# [left + 1, j] 这个区间里的元素都严格小于 pivot
for i in range(left + 1, right + 1):
if input[i] < pivot:
j += 1
input[i], input[j] = input[j], input[i]
input[left], input[j] = input[j], input[left]
return j
if __name__ == '__main__':
input = [9, 14, 1, 16, 19, 13, 12]
k = 4
solution = Solution()
result = solution.getLeastNumbers_Solution(input, k)
print(result)
Python 代码2：最大堆
class Solution(object):
def getLeastNumbers_Solution(self, input, k):
"""
:type input: list[int]
:type k: int
:rtype: list[int]
"""
size = len(input)
if size == 0:
return []
if k == size:
return sorted(input)
import heapq
l = []
for num in input[:k]:
heapq.heappush(l, -num)
for num in input[k:]:
top = l[0]
if top < -num:
heapq.heappushpop(l, -num)
return sorted([-num for num in l])
当然，Python 中的 heapq 直接就有获取最小 个元素的方法。
Python 代码：
class Solution(object):
def getLeastNumbers_Solution(self, input, k):
"""
:type input: list[int]
:type k: int
:rtype: list[int]
"""
size = len(input)
if size == 0:
return []
if k == size:
return sorted(input)
import heapq
heapq.heapify(input)
return sorted(heapq.nsmallest(k, input))
Python 代码：两路 partition 的写法：
class Solution(object):
def getLeastNumbers_Solution(self, input, k):
"""
:type input: list[int]
:type k: int
:rtype: list[int]
"""
size = len(input)
if size == 0:
return []
if k == size:
return sorted(input)
l = 0
r = size - 1
while l <= r:
p = self.__partition(input, l, r)
if p == k - 1:
return sorted(input[:p + 1])
elif p > k - 1:
# 此时 k-1 p
r = p - 1
else:
# 此时 p k-1
l = p + 1
def __partition(self, input, left, right):
# 只有一个数的时候，就没有必要 partition 了
# 直接返回这个数的索引
if left == right:
return left
pivot = input[left]
l = left + 1
r = right
while True:
while l <= right and input[l] <= pivot:
l += 1
while r > left and input[r] >= pivot:
r -= 1
if l > r:
break
input[l], input[r] = input[r], input[l]
l += 1
r -= 1
input[left], input[r] = input[r], input[left]
return r
C++ 代码：
《剑指 Offer （第 2 版）》第 40 题：最小的 K 个数
