数组中第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];
    }
}

Leave a Reply

Your email address will not be published. Required fields are marked *

This site uses Akismet to reduce spam. Learn how your comment data is processed.