如何实现TopK问题的解决方案?之堆排序实现

如何实现TopK问题的解决方案?之堆排序

TopK问题是指在一个数据集合中,找出排名前K的元素。常见的解决方案有两种:堆排序和快速选择。

堆排序

堆排序是一种基于堆的排序算法,它通过维护一个大小为K的小根堆,来找出排名前K的元素。遍历整个数据集合,如果当前元素比堆顶元素大,就将堆顶元素弹出,并将当前元素加入堆中。最终,堆中的元素就是排名前K的元素。

下面是TopK问题的Java实现:

import java.util.PriorityQueue;

public class TopK {
    public static int[] topK(int[] nums, int k) {
        PriorityQueue<Integer> queue = new PriorityQueue<>();
        for (int num : nums) {
            if (queue.size() < k) {
                queue.offer(num);
            } else if (queue.peek() < num) {
                queue.poll();
                queue.offer(num);
            }
        }
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = queue.poll();
        }
        return result;
    }
}