能力随心，自然成长

0%

最小的K个数的多种解法

发表于 更新于 分类于

最小的k个数在面试中遇到了4次，记录一下多种解法。

快速选择

  • 初始版本
    • 时间复杂度最坏情况下O(NlogN)
    • 空间复杂度O(logN)
1
class Solution(object):
2
    def getLeastNumbers(self, arr, k):
3
        """
4
        :type arr: List[int]
5
        :type k: int
6
        :rtype: List[int]
7
        """
8
        if k < 0 or k > len(arr): return
9
        if k == 0: return []
10
        self.quickSort(arr, 0, len(arr) - 1, k)
11
        return arr[:k]
12
13
    def quickSort(self, arr, i, j, k):
14
        if i > j: return
15
        p = self.partition(arr, i, j)
16
        if p == k:
17
            return
18
        elif p > k:
19
            self.quickSort(arr, i, p - 1, k)
20
        else:
21
            self.quickSort(arr, p + 1, j, k)
22
23
    def partition(self, arr, i, j):
24
        # if i >= j: return i
25
        pivot = arr[i]
26
        while i < j:
27
            while i < j and arr[j] >= pivot:
28
                j -= 1
29
            arr[i] = arr[j]
30
31
            while i < j and arr[i] <= pivot:
32
                i += 1
33
            arr[j] = arr[i]
34
35
        arr[i] = pivot
36
        return i

  • 优化版本

    三数取中的方法优化选取枢轴

    1
    def partition(self, arr, i, j):
    2
            mid = (i + j) >> 1
    3
            if arr[i] > arr[j]:  # 把小的放到i上去
    4
                arr[i], arr[j] = arr[j], arr[i]
    5
            if arr[mid] > arr[j]:  # 把最大的放到最后去
    6
                arr[mid], arr[j] = arr[j], arr[mid]
    7
            if arr[mid] > arr[i]:
    8
                arr[mid], arr[i] = arr[i], arr[mid]
    9
    10
            pivot = arr[i]
    11
            while i < j:
    12
                while i < j and arr[j] >= pivot:
    13
                    j -= 1
    14
                arr[i] = arr[j]
    15
    16
                while i < j and arr[i] <= pivot:
    17
                    i += 1
    18
                arr[j] = arr[i]
    19
    20
            arr[i] = pivot
    21
            return i

最大堆

要找到最小的K个数，需要建立一个大顶堆，初始化一个k的元素大顶堆，对于剩下的n-k个数，遍历一下，如果遇到比堆顶元素小的，则删除堆顶元素，该元素入堆，更新堆。

时间复杂度：建堆：O(NlogK)

在python中内置的堆为小顶堆，若要找到最小的K个数，将元素取反即可。

1
class Solution(object):
2
    def getLeastNumbers(self, arr, k):
3
        """
4
        :type arr: List[int]
5
        :type k: int
6
        :rtype: List[int]
7
        """
8
        if k < 0 or k > len(arr): return
9
        if k == 0: return []
10
        import heapq
11
        arr = [-a for a in arr]
12
        hp = arr[:k]
13
        heapq.heapify(hp)
14
15
        for i in range(k, len(arr)):
16
            if arr[i] > hp[0]:
17
                heapq.heappop(hp)
18
                heapq.heappush(hp, arr[i])
19
20
        return [-i for i in hp]

计数排序

数据范围有限时直接计数排序就行了：时间复杂度O(N)

1
class Solution(object):
2
    def getLeastNumbers(self, arr, k):
3
        """
4
        :type arr: List[int]
5
        :type k: int
6
        :rtype: List[int]
7
        """
8
9
        counter = [0] * 10000
10
        for a in arr:
11
            counter[a] += 1
12
13
        ans = []
14
        idx = 0
15
16
        for i in range(len(counter)):
17
            while counter[i] > 0 and idx < k:
18
                ans.append(i)
19
                idx += 1
20
                counter[i] -= 1
21
            if idx == k: break
22
        return ans