剑指29-最小的K个数

Posted on | In |

题目描述

输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4,。

长度为k的辅助数组。

  • 每次遍历到一个数，可能需要给辅助数组排序，因此时间复杂度是\(n\log k\)
1
2
3
4
5
6
7
8
9
10
11
class Solution:
    def GetLeastNumbers_Solution(self, tinput, k):
        if len(tinput) < k or k <=0:
            return []
        lk = tinput[:k]
        lk.sort()
        for i in range(k, len(tinput)):
            if lk[-1] > tinput[i]:
                lk[-1] = tinput[i]
            lk.sort()
        return lk

用内置的堆实现heapq，性能也不错： 

1
2
3
4
5
6
7
8
import heapq
 
class Solution:
    def GetLeastNumbers_Solution(self, tinput, k):
        if not tinput or not k or k > len(tinput):
            return []
        heapq.heapify(tinput)
        return [heapq.heappop(tinput) for _ in xrange(k)]

用快排的partition方法来找到第k大的数，那么该数以及其左侧k-1个数一起，就是最小的k个数。

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
class Solution:
    def GetLeastNumbers_Solution(self, tinput, k):
        def partition(nums, l, r):
            x = nums[r]
            i, j = l, r-1
            while i < j: 
                while nums[i] < x and i < j: # 循环终止时i满足nums[i]>=x或者i==j
                    i += 1
                while nums[j] >= x and i < j: # 循环终止时j满足nums[j]<x或者i==j
                    j -= 1
                if j > i: # 说明必有nums[i]>=x且nums[j]<x，需要调换二者位置
                    nums[i], nums[j] = nums[j], nums[i]

            if nums[i] < x: 
                i += 1
            nums[i], nums[r] = x, nums[i] 
            return i
        if tinput is None or len(tinput) < k or k <= 0:
            return []
        s, t = 0, len(tinput)-1
        while True:
            ix = partition(tinput, s, t)
            if ix == k -1:
                return sorted(tinput[:k])
            elif ix > k -1:
                t = ix -1
            else:
                s = ix + 1