LeetCode 剑指 Offer 40. 最小的k个数

Categories:

Tags:

题目描述
思路分析
1. 解法一：大顶堆（推荐）
2. 解法二：快速选择
相似题目

题目描述

✅ 剑指 Offer 40. 最小的k个数

思路分析

解法一：大顶堆（推荐）

核心思路：

维护一个大小为 k 的大顶堆，堆顶始终是堆中最大的元素

遍历数组，若堆未满则直接入堆；若堆已满且当前元素小于堆顶，则弹出堆顶并将当前元素入堆

遍历结束后堆中保存的就是最小的 k 个数

复杂度分析：

时间复杂度：O(n log k)，其中 n 为数组长度，每次堆操作为 O(log k)

空间复杂度：O(k)，堆的大小为 k

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) {
            return new int[0];
        }
        // 大顶堆，堆顶为堆中最大值
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
        for (int num : arr) {
            // 堆未满时直接加入
            if (maxHeap.size() < k) {
                maxHeap.offer(num);
            } else if (num < maxHeap.peek()) {
                // 当前元素小于堆顶，替换堆顶
                maxHeap.poll();
                maxHeap.offer(num);
            }
        }
        // 将堆中元素转为数组
        int[] res = new int[maxHeap.size()];
        int i = 0;
        for (int val : maxHeap) {
            res[i++] = val;
        }
        return res;
    }
}

// 大顶堆定义
type MaxHeap []int

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *MaxHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func getLeastNumbers(arr []int, k int) []int {
	if k == 0 {
		return []int{}
	}
	h := &MaxHeap{}
	heap.Init(h)
	for _, num := range arr {
		// 堆未满时直接入堆
		if h.Len() < k {
			heap.Push(h, num)
		} else if num < (*h)[0] {
			// 当前元素小于堆顶，弹出堆顶并入堆
			heap.Pop(h)
			heap.Push(h, num)
		}
	}
	// 收集堆中所有元素
	res := make([]int, h.Len())
	for i, v := range *h {
		res[i] = v
	}
	return res
}

解法二：快速选择

核心思路：

基于快速排序的 partition 思想，每次选取一个基准值，将数组分为小于基准和大于基准两部分

若左侧元素个数恰好等于 k，则左侧即为最小的 k 个数

若左侧元素个数小于 k，则递归处理右侧；否则递归处理左侧

平均时间复杂度优于堆，但不适用于数据流场景

复杂度分析：

时间复杂度：O(n)，其中 n 为数组长度，平均情况下快速选择为线性复杂度

空间复杂度：O(log n)，递归调用栈深度

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) {
            return new int[0];
        }
        quickSelect(arr, 0, arr.length - 1, k);
        return Arrays.copyOf(arr, k);
    }

    private void quickSelect(int[] arr, int left, int right, int k) {
        if (left >= right) {
            return;
        }
        // 执行 partition，返回基准值最终位置
        int pivotIdx = partition(arr, left, right);
        if (pivotIdx + 1 == k) {
            // 基准左侧恰好有 k 个元素
            return;
        } else if (pivotIdx + 1 < k) {
            // 最小的 k 个数还在右侧
            quickSelect(arr, pivotIdx + 1, right, k);
        } else {
            // 最小的 k 个数在左侧
            quickSelect(arr, left, pivotIdx - 1, k);
        }
    }

    private int partition(int[] arr, int left, int right) {
        int pivot = arr[left];
        int i = left + 1;
        int j = right;
        while (true) {
            // 从左找第一个大于基准的位置
            while (i <= j && arr[i] <= pivot) {
                i++;
            }
            // 从右找第一个小于等于基准的位置
            while (i <= j && arr[j] > pivot) {
                j--;
            }
            if (i > j) {
                break;
            }
            int tmp = arr[i];
            arr[i] = arr[j];
            arr[j] = tmp;
        }
        // 将基准放到正确位置
        arr[left] = arr[j];
        arr[j] = pivot;
        return j;
    }
}

func getLeastNumbers2(arr []int, k int) []int {
	if k == 0 {
		return []int{}
	}
	quickSelect(arr, 0, len(arr)-1, k)
	return arr[:k]
}

func quickSelect(arr []int, left, right, k int) {
	if left >= right {
		return
	}
	// 执行 partition，返回基准值最终位置
	pivotIdx := partition(arr, left, right)
	if pivotIdx+1 == k {
		return
	} else if pivotIdx+1 < k {
		// 最小的 k 个数还在右侧
		quickSelect(arr, pivotIdx+1, right, k)
	} else {
		// 最小的 k 个数在左侧
		quickSelect(arr, left, pivotIdx-1, k)
	}
}

func partition(arr []int, left, right int) int {
	pivot := arr[left]
	i, j := left+1, right
	for {
		// 从左找第一个大于基准的位置
		for i <= j && arr[i] <= pivot {
			i++
		}
		// 从右找第一个小于等于基准的位置
		for i <= j && arr[j] > pivot {
			j--
		}
		if i > j {
			break
		}
		arr[i], arr[j] = arr[j], arr[i]
	}
	// 将基准放到正确位置
	arr[left], arr[j] = arr[j], arr[left]
	return j
}

相似题目

题目	难度	考察点
215. 数组中的第K个最大元素	中等	堆 / 快速选择
347. 前 K 个高频元素	中等	堆 / 桶排序 / 哈希表
373. 查找和最小的 K 对数字	中等	小顶堆 / 优先队列
378. 有序矩阵中第 K 小的元素	中等	堆 / 二分查找
451. 根据字符出现频率排序	中等	堆 / 桶排序 / 哈希表
692. 前K个高频单词	中等	堆 / 哈希表
703. 数据流中的第 K 大元素	简单	小顶堆 / 优先队列
1046. 最后一块石头的重量	简单	大顶堆 / 优先队列

本文作者: HGNULB

本文链接: https://hgnulb.github.io/blog/2026/15293004