数组中第k大的元素

此文为 LeetCode 215. Kth Largest Element in an Array 的题解。

题意

在一个未排序的数组中，寻找第k大的元素。

分析

这个题目其实很简单，但是踩了一些坑，所以记录一下：

首先是用PriorityQueue，然后poll k次即可（5 ms，37.5 MB）。但是属于API caller，略去不表。

其次就是堆排序，很久不写了，第一次把堆排序写成了选择排序，导致时间很长：101 ms，37.2 MB。

后来复习了下堆排序，写出来了（2 ms，36.6 MB）：

// from https://www.robberphex.com/kth-largest-element-in-an-array/
// 2 ms，36.6 MB
class Solution {
    // 调整current子树，将其整理成最小堆
    private void adjust(int[] nums, int current, int length) {
        // 将堆顶取出来
        int num = nums[current];
        // 将子树中比顶大的数依次弹上来
        for (int i = current * 2 + 1; i < length; i = 2 * i + 1) {
            if (i + 1 < length && nums[i + 1] < nums[i]) {
                i++;
            }
            if (nums[i] < num) {
                nums[current] = nums[i];
                current = i;
            } else {
                break;
            }
        }
        // 不能弹，则将该数放到对应的位置
        nums[current] = num;
    }

    public int findKthLargest(int[] nums, int k) {
        for (int i = (nums.length - 1) / 2; i >= 0; i--) {
            // 从第一个非叶子结点开始，构建堆
            adjust(nums, i, nums.length);
        }
        for (int i = 1; i <= nums.length - k+1; i++) {
            // 将堆顶交换到数组尾部
            int tmp = nums[nums.length - i];
            nums[nums.length - i] = nums[0];
            nums[0] = tmp;
            // 然后继续构建堆
            adjust(nums, 0, nums.length - i);
        }

        return nums[k - 1];
    }
}

当然，很容易想到，最优解就是快速选择算法：快速排序划分区间的时候，忽略所有不包含k的区间即可（1 ms，37.3 MB）：

// from https://www.robberphex.com/kth-largest-element-in-an-array/
// 1 ms,37.3 MB
class Solution {
    public int findKthLargest(int[] nums, int k) {
        int high = nums.length - 1;
        int low = 0;
        while (low < high) {
            // 选择的区间
            int i = low, j = high;
            // 从区间两头向中间推进
            int pivot = nums[i + (j - i) / 2];
            while (i < j) {
                // 跳过符合要求的元素
                while (i < j && nums[i] > pivot) {
                    i++;
                }
                while (i < j && nums[j] < pivot) {
                    j--;
                }
                // 如果元素相等，右侧向左移动一位，确保区间一直在缩小
                // 同时，确保i是区间的中点（即pivot所在的位置
                if (nums[i] == nums[j]) {
                    j--;
                } else {
                    // 交换元素
                    int tmp = nums[i];
                    nums[i] = nums[j];
                    nums[j] = tmp;
                }
            }
            if (i < k - 1) {
                // k在右侧
                low = i + 1;
            } else if (i > k - 1) {
                // k在左侧
                high = i - 1;
            } else {
                // k已经找到
                break;
            }
        }

        return nums[k - 1];
    }
}