题目描述
输入n个整数,找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字,则最小的4个数字是1,2,3,4,。
解题思路
不用动脑的解题思路:实现一个排序算法,然后返回前k个数字。这种思路很low,但是考虑到可以帮自己温习一下快排和归并排序算法,就先用这两种排序算法实现了下:
#快速排序,O(nLogn)
def GetLeastNumbers_Solution(self, tinput, k):
if len(tinput) < k:
return []
return self.quick_sort(tinput)[:k]
def quick_sort(self, list):
if len(list) < 2:
return list[:]
left = (self.quick_sort([i for i in list[1:] if i <= list[0]]))
right = (self.quick_sort([i for i in list[1:] if i > list[0]]))
return left + [list[0]] + right
#归并排序,O(nLogn)
def GetLeastNumbers_Solution(self, tinput, k):
if len(tinput) < k:
return []
return self.merge_sort(tinput)[:k]
def merge_sort(self, list):
if len(list) < 2:
return list[:]
left = self.merge_sort(list[:len(list)//2])
right = self.merge_sort(list[len(list)//2:])
return self.merge(left,right)
def merge(self,left, right):
res = []
while left and right:
res.append(left.pop(0)) if left[0] < right[0] else res.append(right.pop(0))
res += left if not right else right
return res
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
- 19
- 20
- 21
- 22
- 23
- 24
- 25
- 26
- 27
- 28
- 29
- 30
- 31
- 32
基于“先整体排序,然后取前k个元素”的思路引申:需要的是前k小的数,没必要舍近求远把整个数组都排好序。这样一来,时间复杂度可以降到O(nLogk)。利用最大堆算法,只不过在这里稍稍将最大堆算法变为最小堆算法,代码如下:
# 最大堆,O(nLogk)
def GetLeastNumbers_Solution(self, tinput, k):
if len(tinput) < k:
return []
res = []
for i in tinput:
heapq.heappush(res, -i) if len(res) < k else heapq.heappushpop(res, -i)
return sorted(list(map(lambda x: -x, res)))
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
番外:Python中的heapq
模块用来建立“堆”这种数据结构。
heapq.heappush(res, -i)
意为:向堆res
中添加一个元素-i
heapq.heappushpop(res, -i)
意为:将元素-i
与堆顶的元素比较。如果该元素值大于堆顶元素,则将该元素与堆顶元素替换。否则不改变堆元素。
解题方法
这个是面试被问过的题目,做法就是使用一个堆来做,特别适合当数据很大的时候。这个其实相当于一个信息检索的过程。据说百度喜欢问此类问题。
heapq.nlargest(n, iterable, key=None) 返回最大的n个元素(Top-K问题)
heapq.nsmallest(n, iterable, key=None) 返回最小的n个元素(Top-K问题)
- 1
- 2
- 3
代码:
# -*- coding:utf-8 -*-
import heapq
class Solution:
def GetLeastNumbers_Solution(self, tinput, k):
if len(tinput) < k:
return []
return heapq.nsmallest(k, tinput)
- 1
- 2
- 3
- 4
- 5
- 6
- 7