算法 - 数组第 K 大元素 / TopK 问题

Posted on

1. 题目

给定整数数组 nums 和整数 k，返回数组中第 k 大的元素（排序后第 k 个位置的值，不是第 k 个不同元素）。

示例:

1
2
3
4
5
输入: nums = [3,2,1,5,6,4], k = 2
输出: 5

输入: nums = [3,2,3,1,2,4,5,5,6], k = 4
输出: 4

2. 解题思路

「第 K 大 / Top K」是超高频题型，有三种主流套路，各有适用场景。

2.1 排序法 O(n log n)

直接排序后取 nums[n - k]。简单，但不是最优，面试可作为保底并说明「有更快的」。

2.2 大小为 K 的最小堆 O(n log k) —— 求 TopK 首选

维护一个只装 k 个元素的最小堆：

  • 遍历数组，堆大小 < k 就入堆。
  • 否则若当前值 > 堆顶（堆里最小的），弹出堆顶、压入当前值。
  • 遍历结束，堆顶就是第 k 大（堆里 k 个数中最小的那个，恰为整体第 k 大）。

优势：天然适合数据流 / 海量数据，内存只需 O(k)，不必一次性拿到全部数据，这也是「TopK」在工程里（如热点搜索词）的常见解法。

2.3 快速选择（Quickselect）平均 O(n)

复用快排的 partition：每趟把基准放到最终位置 p，若 p === n-k 直接返回；否则只递归目标所在的那一半，平均线性、最坏 O(n²)（随机化 pivot 可高概率避免）。

  • 时间复杂度：平均 O(n)；空间 O(1)（原地）。

3. TypeScript 实现

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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
// 2.3 快速选择（原地，平均 O(n)）
function findKthLargest(nums: number[], k: number): number {
    const target = nums.length - k; // 转成升序下标

    function partition(lo: number, hi: number): number {
        // 随机化，避免有序退化
        const rand = lo + Math.floor(Math.random() * (hi - lo + 1));
        [nums[rand], nums[hi]] = [nums[hi], nums[rand]];
        const pivot = nums[hi];
        let i = lo;
        for (let j = lo; j < hi; j++) {
            if (nums[j] < pivot) {
                [nums[i], nums[j]] = [nums[j], nums[i]];
                i++;
            }
        }
        [nums[i], nums[hi]] = [nums[hi], nums[i]];
        return i;
    }

    let lo = 0;
    let hi = nums.length - 1;
    while (true) {
        const p = partition(lo, hi);
        if (p === target) return nums[p];
        if (p < target) lo = p + 1;
        else hi = p - 1;
    }
}

// 2.2 最小堆版（适合数据流 / 海量 TopK）
// 借助一个数值小顶堆：这里用极简实现演示思路
class MinHeap {
    private h: number[] = [];
    get size() { return this.h.length; }
    peek() { return this.h[0]; }
    push(v: number) {
        const h = this.h;
        h.push(v);
        let i = h.length - 1;
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (h[p] <= h[i]) break;
            [h[p], h[i]] = [h[i], h[p]];
            i = p;
        }
    }
    pop() {
        const h = this.h;
        const top = h[0];
        const last = h.pop()!;
        if (h.length) {
            h[0] = last;
            let i = 0;
            while (true) {
                const l = 2 * i + 1;
                const r = 2 * i + 2;
                let s = i;
                if (l < h.length && h[l] < h[s]) s = l;
                if (r < h.length && h[r] < h[s]) s = r;
                if (s === i) break;
                [h[s], h[i]] = [h[i], h[s]];
                i = s;
            }
        }
        return top;
    }
}

function topKKthLargest(nums: number[], k: number): number {
    const heap = new MinHeap();
    for (const x of nums) {
        if (heap.size < k) heap.push(x);
        else if (x > heap.peek()) {
            heap.pop();
            heap.push(x);
        }
    }
    return heap.peek();
}

console.log(findKthLargest([3, 2, 1, 5, 6, 4], 2));     // 5
console.log(topKKthLargest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)); // 4

4. 面试延伸

  • 前 K 个高频元素（347）：先哈希统计频次，再对频次跑「大小 k 的小顶堆」或快速选择，是 TopK + 堆的经典组合。
  • 海量数据（如 10 亿数取前 100）：绝不可能全放内存，答案就是大小为 k 的小顶堆分块处理，凸显堆相对于快选的优势。
  • 第 K 小的元素 / 有序矩阵第 K 小（378）：可加一条「二分答案」解法——在值域二分，用 O(n) 统计 <= mid 的个数，O(n log(range))。
  • 一定要区分：快速选择平均 O(n) 但会打乱原数组；堆法保证 O(n log k) 且能处理流式数据。选型看是否可修改原数组、是否一次性数据。

难度：中等 | LeetCode 215 题 | 堆 / 分治双高频