TOP-K 算法详解：最小堆（Min-Heap）解法#

1
import java.util.*;
2

3
/**
4
 * 使用最小堆实现的 Top-K 算法
5
 *
6
 * @param <K> 可比较的元素类型
7
 */
8
public class TopKHeap<K> {
9
    private final int k;                              // K 值
10
    private final Comparator<? super K> comparator;   // 比较器
11
    private final List<K> heap;                       // 堆的存储
12

13
    /**
14
     * 构造函数
15
     *
16
     * @param k Top-K 中的 K 值
17
     * @param comparator 元素比较器
18
     */
19
    public TopKHeap(int k, Comparator<? super K> comparator) {
20
        if (k <= 0) {
21
            throw new IllegalArgumentException("k must be positive");
22
        }
23
        this.k = k;
24
        this.comparator = comparator;
25
        this.heap = new ArrayList<>(k);
26
    }
27

28
    /**
29
     * 添加一个元素
30
     *
31
     * @param element 待添加的元素
32
     */
33
    public void offer(K element) {
34
        if (element == null) {
35
            return;
36
        }
37

38
        // 情况1：堆未满，直接插入并向上调整
39
        if (heap.size() < k) {
40
            heap.add(element);
41
            siftUp(heap.size() - 1);
42
        }
43
        // 情况2：堆已满，但新元素比堆顶大，替换并向下调整
44
        else if (comparator.compare(element, heap.get(0)) > 0) {
45
            heap.set(0, element);
46
            siftDown(0);
47
        }
48
        // 情况3：新元素小于等于堆顶，无需处理
49
    }
50

51
    /**
52
     * 批量添加元素
53
     *
54
     * @param elements 待添加的元素集合
55
     */
56
    public void addAll(Collection<K> elements) {
57
        for (K element : elements) {
58
            offer(element);
59
        }
60
    }
61

62
    /**
63
     * 获取 Top-K 结果
64
     * 返回一个新的列表，避免外部修改
65
     *
66
     * @return Top-K 元素列表
67
     */
68
    public List<K> getResult() {
69
        return new ArrayList<>(heap);
70
    }
71

72
    /**
73
     * 获取当前堆中元素数量
74
     */
75
    public int size() {
76
        return heap.size();
77
    }
78

79
    /**
80
     * 判断是否包含指定元素
81
     */
82
    public boolean contains(K element) {
83
        return heap.contains(element);
84
    }
85

86
    /**
87
     * 向上调整（sift-up）
88
     * 用于插入新元素时恢复堆性质
89
     *
90
     * 时间复杂度：O(log K)
91
     */
92
    private void siftUp(int index) {
93
        while (index > 0) {
94
            int parent = (index - 1) / 2;
95

96
            if (comparator.compare(heap.get(index), heap.get(parent)) < 0) {
97
                swap(index, parent);
98
                index = parent;
99
            } else {
100
                break;
101
            }
102
        }
103
    }
104

105
    /**
106
     * 向下调整（sift-down）
107
     * 用于删除堆顶或替换堆顶后恢复堆性质
108
     *
109
     * 时间复杂度：O(log K)
110
     */
111
    private void siftDown(int index) {
112
        while (true) {
113
            int smallest = index;
114
            int left = 2 * index + 1;
115
            int right = 2 * index + 2;
116

117
            // 与左子节点比较
118
            if (left < heap.size() &&
119
                comparator.compare(heap.get(left), heap.get(smallest)) < 0) {
120
                smallest = left;
121
            }
122

123
            // 与右子节点比较
124
            if (right < heap.size() &&
125
                comparator.compare(heap.get(right), heap.get(smallest)) < 0) {
126
                smallest = right;
127
            }
128

129
            // 如果最小值不是当前节点，交换并继续向下调整
130
            if (smallest != index) {
131
                swap(index, smallest);
132
                index = smallest;
133
            } else {
134
                break;
135
            }
136
        }
137
    }
138

139
    /**
140
     * 交换堆中两个位置的元素
141
     */
142
    private void swap(int i, int j) {
143
        K temp = heap.get(i);
144
        heap.set(i, heap.get(j));
145
        heap.set(j, temp);
146
    }
147
}

1
public class Main {
2
    public static void main(String[] args) {
3
        int[] nums = {9, 3, 7, 1, 5, 8, 2, 6, 4};
4
        int k = 3;
5

6
        TopKHeap<Integer> topK = new TopKHeap<>(k, Integer::compareTo);
7

8
        System.out.println("=== Top-K 元素查找过程 ===");
9
        System.out.println("输入数组: " + Arrays.toString(nums));
10
        System.out.println("K = " + k);
11
        System.out.println();
12

13
        for (int num : nums) {
14
            System.out.println("处理元素: " + num);
15
            topK.offer(num);
16
            System.out.println("当前堆: " + topK.getResult());
17
        }
18

19
        System.out.println();
20
        System.out.println("=== 最终结果 ===");
21
        System.out.println("Top-" + k + ": " + topK.getResult());
22

23
        // 验证结果
24
        System.out.println();
25
        System.out.println("=== 验证 ===");
26
        Integer[] sorted = Arrays.stream(nums)
27
                                  .boxed()
28
                                  .sorted(Comparator.reverseOrder())
29
                                  .toArray(Integer[]::new);
30
        List<Integer> expected = Arrays.asList(sorted).subList(0, k);
31
        System.out.println("期望结果: " + expected);
32
        System.out.println("验证通过: " + topK.getResult().containsAll(expected));
33
    }
34
}

1
=== Top-K 元素查找过程 ===
2
输入数组: [9, 3, 7, 1, 5, 8, 2, 6, 4]
3
K = 3
4

5
处理元素: 9
6
当前堆: [9]
7
处理元素: 3
8
当前堆: [3, 9]
9
处理元素: 7
10
当前堆: [3, 9, 7]
11
处理元素: 1
12
当前堆: [1, 9, 7]      (1 < 3，跳过，堆结构不变但堆顶不变)
13
处理元素: 5
14
当前堆: [5, 9, 7]      (5 > 3，替换堆顶3，调整后: [5, 9, 7])
15
处理元素: 8
16
当前堆: [7, 9, 8]      (8 > 5，替换堆顶5，调整后: [7, 9, 8])
17
处理元素: 2
18
当前堆: [2, 9, 7]      (2 < 7，跳过)
19
处理元素: 6
20
当前堆: [6, 9, 7]      (6 > 2，替换堆顶2，调整后: [6, 9, 7])
21
处理元素: 4
22
当前堆: [4, 9, 7]      (4 > 2，替换堆顶2，调整后: [4, 9, 7])
23

24
=== 最终结果 ===
25
Top-3: [4, 9, 7]
26

27
=== 验证 ===
28
期望结果: [9, 8, 7]
29
验证通过: true

1
输入: [9, 3, 7, 1, 5, 8, 2, 6, 4], K = 3
2

3
Step 1: 处理 9
4
  ┌─────┐
5
  │  9  │
6
  └─────┘
7
  堆: [9]
8

9
Step 2: 处理 3
10
       3
11
      /
12
     9
13
  堆: [3, 9]
14

15
Step 3: 处理 7
16
       3
17
      / \
18
     9   7
19
  堆: [3, 9, 7]     (堆已满，size = K)
20

21
Step 4: 处理 1
22
  ┌─────────────────────────────┐
23
  │ 1 < 3 (堆顶)，跳过          │
24
  └─────────────────────────────┘
25
  堆: [3, 9, 7]     (不变)
26

27
Step 5: 处理 5
28
  ┌─────────────────────────────┐
29
  │ 5 > 3 (堆顶)，替换并调整    │
30
  └─────────────────────────────┘
31

32
  替换: [5, 9, 7]
33
        ↓
34
    小顶堆化 (sift-down)
35
        ↓
36
  堆: [5, 9, 7]
37

38
Step 6: 处理 8
39
  ┌─────────────────────────────┐
40
  │ 8 > 5 (堆顶)，替换并调整    │
41
  └─────────────────────────────┘
42

43
  替换: [8, 9, 7]
44
        ↓
45
    小顶堆化
46
        ↓
47
  堆: [7, 9, 8]
48

49
Step 7-9: 处理 2, 6, 4
50
  2 < 7，跳过
51
  6 < 7，跳过
52
  4 < 7，跳过
53

54
───────────────────────────────
55
最终堆: [7, 9, 8]
56
Top-3: [7, 8, 9] ✓

1
import java.util.*;
2

3
public class TopKHeapDemo {
4
    public static void main(String[] args) {
5
        // 示例：找出数组中最大的 3 个元素
6
        int[] nums = {9, 3, 7, 1, 5, 8, 2, 6, 4};
7
        int k = 3;
8

9
        TopKHeap<Integer> topK = new TopKHeap<>(k, Integer::compareTo);
10

11
        for (int num : nums) {
12
            topK.offer(num);
13
        }
14

15
        System.out.println("Top-" + k + ": " + topK.getResult());
16
        System.out.println("复杂度: O(N log K) 时间, O(K) 空间");
17
    }
18
}