最小的K个数

最小的K个数_牛客题霸_牛客网

问题简述

给定一个长度为 n 的可能有重复值的数组，找出其中不去重的最小的 k 个数。例如数组元素是4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4(任意顺序皆可)。

思路1：快排（推荐）

Python

class Solution:
    def GetLeastNumbers_Solution(self, tinput, k):
        
        def partition(a, lo, hi):
            if lo >= hi: return 
            
            p = a[lo]
            l, r = lo, hi
            while l < r:
                while l < r and a[r] >= p: r -= 1
                while l < r and a[l] <= p: l += 1
                a[l], a[r] = a[r], a[l]
            
            a[lo], a[l] = a[l], a[lo]
            
            # 因为只需要前 k 个数，所有加上 if，去掉 if 就是标准的快排
            if l > k - 1: partition(a, lo, l - 1)
            if l < k - 1: partition(a, l + 1, hi)
        
        partition(tinput, 0, len(tinput) - 1)
        return tinput[: k]

思路2：堆排

• O(nlogk) 的写法（推荐）：

  • 取前 k 个数建大顶堆；

  • 然后遍历剩余的数，如果小于堆顶就 pushpop

• O(nlogn) 的写法：

  • 整体构建小顶堆，然后弹出前 k 个；

Python

class Solution:
    def GetLeastNumbers_Solution(self, a, k):
        if k == 0: return []
        
        import heapq
        
        heapq.heapify(h := [-a[i] for i in range(k)])  # 取相反数构建大顶堆
        for i in range(k, len(a)):
            if a[i] < -h[0]:
                heapq.heappushpop(h, -a[i])
        
        return [-x for x in h]