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

279 阅读2分钟

方法一:快速选择,O(N)

  • 理想来说,partition应该刚刚把数组平分,因此总遍历数是n+n/2+n/4+...+1=o(n),
  • 最差情况是数组有序,O(N^2)
  • 判断左右指针相遇的位置的index是不是length - k
    • 如果是,那就返回
    • 相遇位置大于length - k,往前找
    • 相遇位置小于length - k,往后找
class Solution {
    public int findKthLargest(int[] nums, int k) {
        sort(nums, 0, nums.length - 1, k);
        return nums[nums.length - k];
    }
    public void sort(int[] nums, int left, int right, int k) {
        if (left > right) {
            return;
        }
        int rnd = new Random().nextInt(right - left + 1) + left;
        swap(nums, left, rnd);
        int pivot = left;
        int i = left, j = right;
        while (i != j) {
            while (nums[j] > nums[pivot]) {//这里写>=可能会出现-1越界
                j--;
            }
            while (i < j &&nums[i] <= nums[pivot]) {
                i++;
            }
            swap(nums, i, j);
        }
        //此时ij相遇
        swap(nums, i, pivot);
        //如果相遇位置正好为倒数第K个,那么就是第K个最大元素的位置
        if (i == nums.length - k) {
            return;
        }
        //比快排省去了一半的比较
        if (i > nums.length - k) {
            sort(nums, left, i - 1, k);
        }
        //比快排省去了一半的比较
        if (i < nums.length - k) {
            sort(nums, i + 1, right, k);
        }
    }
    public void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}

方法二:小顶堆,O(NlogK)

维护一个含有K个元素的小顶堆,最终堆顶即为第K个最大元素。

  • 以数组的前K个元素建一个小顶堆。
  • 扫一遍数组的剩下N-K个元素,当遇到大于堆顶的元素,删掉堆顶,把这个元素加入堆并build。
  • 扫完后,堆顶即为结果。

1. 手写小顶堆

//自己实现小顶堆的build,O(NlogK)
class Solution {
    public int findKthLargest(int[] nums, int k) {
        buildMinHeap(nums, k); //构造大小为K的小顶堆
        //从第K+1个数开始遍历
        for (int i = k; i < nums.length; i++) {
            if (nums[i] > nums[0]) {
                swap(nums, 0, i);
                heapify(nums, 0, k); 
            }
        }
        return nums[0];
    }
    
    // n为heapify的索引上限,不包括
    public void heapify(int[] nums, int i, int n) {
        int c1 = i * 2 + 1;
        int c2 = i * 2 + 2;
        int min = i;
        if (c1 < n && nums[c1] < nums[min]) {
            min = c1;
        }
        if (c2 < n && nums[c2] < nums[min]) {
            min = c2;
        }
        if (min != i) {
            swap(nums, min, i);
            heapify(nums, min, n);
        }
    }

    // 构造小顶堆,k为堆内元素个数
    public void buildMinHeap(int[] nums, int k) {
        int lastParentIndex = (k - 2) / 2; //最后一个父节点的索引
        // 把前k个数构造成小顶堆
        for (int i = lastParentIndex; i >= 0; i--) {
            heapify(nums, i, k);
        }
    }

    public void swap(int[] nums, int m, int n) {
        int tmp = nums[m];
        nums[m] = nums[n];
        nums[n] = tmp;
    }
}

2.Java的PriorityQueue当作小顶堆

class Solution {
    public int findKthLargest(int[] nums, int k) {
        Queue<Integer> minHeap = new PriorityQueue<>();
        for (int x : nums) {
            if (minHeap.size() < k) {
                minHeap.add(x);
            } else if (x > minHeap.peek()) {
                minHeap.poll();
                minHeap.add(x);
            }
        }
        return minHeap.peek();
    }
}

方法三:大顶堆,O(NlogN)

//大顶堆O(NlogN)
class Solution {
    public int findKthLargest(int[] nums, int k) {
        buildMaxHeap(nums);
        int i = 0;
        //一共换K次,直到把第K大的元素换上来
        for (i = nums.length - 1; i >= nums.length - k; i--) {
            swap(nums, 0, i);
            maxHeapify(nums, i , 0);
        }
        return nums[++i];//到了倒数第k个后又--了,所以得++

    }
    public void maxHeapify(int[] nums, int n, int i) {
        int c1 = i * 2 + 1;
        int c2 = i * 2 + 2;
        int max = i;
        if (c1 < n && nums[c1] > nums[max]) {
            max = c1;
        }
        if (c2 < n && nums[c2] > nums[max]) {
            max = c2;
        }
        if (max != i) {
            swap(nums, max, i);
            maxHeapify(nums, n, max);
        }
    }
    public void buildMaxHeap(int[] nums) {
        int lastNode = nums.length - 1;
        int lastParent = (lastNode - 1) / 2;
        for (int i = lastParent; i >= 0; i--) {
            maxHeapify(nums, nums.length, i);
        }
    }
    public void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}

升华

  • topk (前k大)用小根堆,维护堆大小不超过 k 即可。每次压入堆前和堆顶元素比较,如果比堆顶元素还小,直接扔掉,否则压入堆。检查堆大小是否超过 k,如果超过,弹出堆顶。复杂度是 nlogk。避免使用大根堆,因为你得把所有元素压入堆,复杂度是 nlogn,而且还浪费内存。如果是海量元素,那就挂了。

  • 求前 k 大,用小根堆,求前 k 小,用大根堆。