29 最小的K个数

题目描述

输入n个整数，找出其中最小的K个数。

示例：

输入: arr = [4,5,1,6,2,7,3,8], k = 4
输出: [1,2,3,4]

解题思路

方法一：排序

核心思想：直接排序后取前k个。

方法二：大根堆

核心思想：维护一个大小为k的大根堆，遍历数组：

• 如果堆未满，加入元素
• 如果元素小于堆顶，替换堆顶

方法三：快速选择（Partition）

核心思想：利用快速排序的partition思想。

时间复杂度：O(n log k)（堆）或 O(n)（快速选择） 空间复杂度：O(k)（堆）或 O(1)（快速选择）

Python 实现

# 方法一：排序
def get_least_numbers_sort(arr, k):
    """
    排序法
    """
    if not arr or k <= 0:
        return []
    return sorted(arr)[:k]


# 方法二：大根堆
import heapq

def get_least_numbers_heap(arr, k):
    """
    大根堆法（使用Python的小根堆模拟）
    """
    if not arr or k <= 0:
        return []

    # Python的heapq是小根堆，存储负数模拟大根堆
    max_heap = []

    for num in arr:
        if len(max_heap) < k:
            heapq.heappush(max_heap, -num)
        elif num < -max_heap[0]:
            heapq.heapreplace(max_heap, -num)

    return [-x for x in sorted(max_heap)]


# 方法三：快速选择
def get_least_numbers_quick_select(arr, k):
    """
    快速选择法
    """
    if not arr or k <= 0:
        return []

    def partition(left, right, pivot_idx):
        pivot = arr[pivot_idx]
        arr[pivot_idx], arr[right] = arr[right], arr[pivot_idx]

        store = left
        for i in range(left, right):
            if arr[i] < pivot:
                arr[store], arr[i] = arr[i], arr[store]
                store += 1

        arr[store], arr[right] = arr[right], arr[(store]
        return store

    def select(left, right, k_smallest):
        if left == right:
            return left

        pivot_idx = partition(left, right, right)

        if k_smallest == pivot_idx:
            return k_smallest
        elif k_smallest < pivot_idx:
            return select(left, pivot_idx - 1, k_smallest)
        else:
            return select(pivot_idx + 1, right, k_smallest)

    select(0, len(arr) - 1, k - 1)
    return sorted(arr[:k])


# 测试代码
if __name__ == "__main__":
    arr = [4, 5, 1, 6, 2, 7, 3, 8]
    k = 4

    print(f"原数组: {arr}")
    print(f"最小的{k}个数（排序）: {get_least_numbers_sort(arr[:], k)}")
    print(f"最小的{k}个数（堆）: {get_least_numbers_heap(arr[:], k)}")

# 输出:
# 原数组: [4, 5, 1, 6, 2, 7, 3, 8]
# 最小的4个数（排序）: [1, 2, 3, 4]
# 最小的4个数（堆）: [1, 2, 3, 4]

图解

大根堆法示例: [4,5,1,6,2,7,3,8], k=4

构建大根堆:
4: 堆=[4]
5: 堆=[5,4]
1: 堆=[5,4,1]
6: 堆=[6,5,1,4]

遍历剩余元素:
2: 2<6, 替换 → 堆=[5,4,1,2]
7: 7>5, 跳过
3: 3<5, 替换 → 堆=[4,3,1,2]
8: 8>4, 跳过

最终堆: [4,3,1,2]
排序后: [1,2,3,4]