数组中的第 K 个最大元素

215. 数组中的第K个最大元素

方法一、基于快速排序的选择算法,时间复杂度 O(n)

这道题是经典的 Top K 问题,可以使用快速排序中的 Partition 思想来解决,而不需要对整个数组进行完全排序。这种方法被称为快速选择(Quickselect)算法。主要思想是:

  1. 选择基准:与快速排序一样,选择一个基准元素(随机选择或最后一个元素)

  2. Partition:将数组进行分成两个部分,左侧都是小于等于基准的元素,右侧都是大于基准的元素

  3. 确认目标位置:

    • 设划分后基准元素的索引是 p

    • 如果 p+1=k,那么 nums[p] 就是第 k 大的元素;

    • 如果 p+1>k,那么第 k 大的元素就在左侧子数组中,递归查找左子数组;

    • 如果 p+1<k,那么第 k 大的元素就在右侧子数组中,递归查找右子数组;

final Random random = new Random();

public int findKthLargest(int[] nums, int k) {
    int n = nums.length;
    // 转换为第 n - k + 1 小
    return quickSelect(nums, 0, n - 1, n - k + 1);
}

public int quickSelect(int[] nums, int left, int right, int k) {
    if (left == right) {
        return nums[left];
    }

    // 随机第一个数为基准
    int pivot = nums[left];

    // 双向扫描
    int i = left - 1, j = right + 1;
    while (i < j) {
        do {
            i++;
        } while (nums[i] < pivot);
        do {
            j--;
        } while (nums[j] > pivot);
        if (i < j) {
            swap(nums, i, j);
        }
    }

    if (j + 1 >= k) {
        return quickSelect(nums, left, j, k);
    } else {
        return quickSelect(nums, j + 1, right, k);
    }
}

public void swap(int[] nums, int i, int j) {
    int tmp = nums[i];
    nums[i] = nums[j];
    nums[j] = tmp;
}

方法二、基于堆排序的选择方法,时间复杂度 O(n\log n)

维护一个最大堆,进行 k-1 次删除,堆顶元素就是我们的答案,下面是使用系统库实现:

public int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> pq = new PriorityQueue<>();
    for (int x : nums){
        pq.offer(-x);
    }
    for (int i = 0; i < k - 1; i++){
        pq.poll();
    }
    return -pq.peek();
}

但是在面试中,面试官可能让我们手写最大堆/最小堆,下面是手动实现最大堆的实现:

class Solution {
    
    int[] nums;

    public int findKthLargest(int[] nums, int k) {
        this.nums = nums;
        int n = nums.length;
        heapify(n);
        for (int i = n - 1; i >= n - k + 1; i--){
            swap(0, i);
            siftDown(0, i);
        }
        return nums[0];
    }

    public void heapify(int heapSize){
        for (int i = parent(heapSize - 1); i >= 0; i--){
            siftDown(i, heapSize);
        }
    }

    public void siftDown(int i, int heapSize){
        int l = 2 * i + 1;
        int r = 2 * i + 2;
        int maxIdx = i;
        if (l < heapSize && nums[l] > nums[maxIdx]){
            maxIdx = l;
        }
        if (r < heapSize && nums[r] > nums[maxIdx]){
            maxIdx = r;
        }
        if (maxIdx == i){
            return;
        }
        swap(maxIdx, i);
        siftDown(maxIdx, heapSize);
    }

    public int parent(int i){
        return (i - 1) / 2;
    }

    public void swap(int i, int j){
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}