【Leetcode 热题 100】215. 数组中的第K个最大元素

问题背景

给定整数数组 n u m s nums nums 和整数 k k k，请返回数组中第 k k k 个最大的元素。

请注意，你需要找的是数组排序后的第 k k k 个最大的元素，而不是第 k k k 个不同的元素。

你必须设计并实现时间复杂度为 O ( n ) O(n) O(n) 的算法解决此问题。

数据约束

1 ≤ k ≤ n u m s . l e n g t h ≤ 1 0 5 1 \le k \le nums.length \le 10 ^ 5 1≤k≤nums.length≤105
− 1 0 4 ≤ n u m s [ i ] ≤ 1 0 4 -10 ^ 4 \le nums[i] \le 10 ^ 4 −104≤nums[i]≤104

解题过程

经典求第 k k k大，最佳方案是用随机选择算法，时间复杂度大致是 O ( N ) O(N) O(N) 这个量级。

考虑到题目分类在堆，那除了调用 API 快速处理的方法之外，再写一下手撕堆当作练习。

相关详细的介绍，可以参考随机快排与随机选择以及堆排序与堆操作。

具体实现

使用堆 API

java 复制代码

class Solution {
    public int findKthLargest(int[] nums, int k) {
        Queue<Integer> queue = new PriorityQueue<>();
        for(int num : nums) {
            queue.offer(num);
        }
        for(int i = 0; i < nums.length - k; i++) {
            queue.poll();
        }
        return queue.peek();
    }
}

自己实现堆

java 复制代码

class Solution {
    public int findKthLargest(int[] nums, int k) {
        int n = nums.length;
        buildHeap(nums);
        // 注意这里要修正索引
        while(k - 1 > 0) {
            swap(nums, 0, --n);
            downAdjust(nums, 0, n);
            k--;
        }
        return nums[0];
    }

    // 交换数组中的两个元素
    private void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }
    
    // 由上而下调整元素
    private void downAdjust(int[] arr, int cur, int size) {
        // 数组下标从零开始，当前节点的左孩子下标为 2 * cur + 1
        int child = 2 * cur + 1;
        while (child < size) {
            // 如果当前节点有右孩子，那么比较两个子节点的值确定潜在的交换对象
            int target = child + 1 < size && arr[child + 1] > arr[child] ? child + 1 : child;
            // 再与当前节点比较大小
            target = arr[target] > arr[cur] ? target : cur;
            // 一旦发现此次操作中无需交换，立即停止流程
            if (target == cur) {
                break;
            }
            // 交换父子节点
            swap(arr, target, cur);
            // 移动工作指针
            cur = target;
            child = 2 * cur + 1;
        }
    }
    
    // 自底向顶建堆
    private void buildHeap(int[] arr) {
        int n = arr.length;
        for (int i = n - 1; i >= 0; i--) {
            downAdjust(arr, i, n);
        }
    }
}

随机选择

java 复制代码

class Solution {
    // 全局变量 first 表示等于当前元素的第一个下标位置，last 表示最后一个
    private static int first, last;
    
    public int findKthLargest(int[] nums, int k) {
        return randomizedSelect(nums, nums.length - k);
    }

    // 交换数组中的两个元素
    private void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }

    private void partition(int[] arr, int left, int right, int pivot) {
        // 初始化维护的指针
        first = left;
        last = right;
        int cur = left;
        // 注意循环到工作指针超越 last 即可结束流程
        while (cur <= last) {
            if (arr[cur] == pivot) {
                cur++; // 当前元素等于划分标准，不作处理并后移指针
            } else if (arr[cur] < pivot) {
                swap(arr, first++, cur++); // 小于划分标准，交换且分别后移两个指针
            } else {
                swap(arr, cur, last--); // 大于划分标准，交换并前移记录最后一个当前元素的指针
            }
        }
    }
    
    private int randomizedSelect(int[] arr, int target) {
        int res = 0;
        for (int left = 0, right = arr.length - 1; left <= right;) {
            partition(arr, left, right, arr[left + (int) (Math.random() * (right - left + 1))]);
            if (target < first) {
                right = first - 1; // 当前下标小于维护的第一个下标，到左边区间中找
            } else if (target > last) {
                left = last + 1; // 当前下标大于维护的最后一个下标，到右边区间中找
            } else {
                res = arr[target]; // 当前下标在维护的范围内，记录结果
                break;
            }
        }
        return res;
    }
}