703. Kth Largest Element in a Stream

Leetcode Heap

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Your KthLargest class will have a constructor which accepts an integer k and an integer array nums, which contains initial elements from the stream. For each call to the method KthLargest.add, return the element representing the kth largest element in the stream.

Example:

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

Note: You may assume that nums' length $≥ k-1$ and $k ≥ 1$.

分析

首先最简单的方法是，对整个数组进行排序，然后通过数组下标索引并返回该元素。时间复杂度是$O(n\log n)$，空间复杂度是$O(1)$.

public int findKthLargest(int[] nums, int k) {
    Arrays.sort(nums);
    return nums[nums.length - k];
}

其次，很自然的想到使用二叉堆处理最大值。在Java中二叉堆可以用PriorityQueue来表示，首先将数组元素依次加入到二叉堆中，然后连续取$k$次最大值即可，第$k$次的返回结果就是第$k$大的值。时间复杂度是$O(n\log n)$，空间复杂度是$O(n)$.

public int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> heap = new PriorityQueue<Integer>(
        nums.length, Collections.reverseOrder()); // 注意堆的顺序reverse         
    heap.addAll(Arrays.asList(num));            // 加入所有元素到堆中
    for (int i = 0; i < k - 1; i ++)
        heap.poll();
    return heap.poll();
}

但是以上的两种方法都不能应用，原因就在于除了findKthLargest方法，还需要实现add方法来添加元素。如果使用方法一，则每一次添加都需要重新排序；如果使用方法二，则每一次添加都需要重新加入所有元素到二叉堆中。算法的时间复杂度达到了$O(n^2)$以上。

所以必须改进以上两种方法。首先比较简单的，改进二叉堆：始终维持二叉堆的大小为$k$，当二叉堆的大小超过$k$时，删除最小值。时间复杂度是$O(n\log k)$，空间复杂度是$O(n)$.

private PriorityQueue<Integer> hp;
private int k;
public KthLargest(int k, int[] nums) {
    this.k = k;
    hp = new PriorityQueue<>();  // 最小二叉堆
    for (int num : nums) {
        hp.offer(num);  // 加入元素
    }
}

public int add(int val) {
    hp.offer(val);
    if (hp.size() > k) hp.poll();
    return hp.peek();
}

注意要先加入元素到二叉堆中，然后再删除最小值。下面方法是错误的，原因在于，当前元素有可能比堆中所有元素要小，这时理应删除该元素，但是却删除了原来堆中的元素。

if (hp.size() >= k) hp.poll();  // 删除最小值
hp.offer(num);                  // 加入元素

其次改进排序的方法：快速选择(quick select)算法，线性时间复杂度！

public int findKthLargest(int[] nums, int k) {
    k = nums.length - k;
    int lo = 0;
    int hi = nums.length - 1;
    while (lo < hi) {
        final int j = partition(nums, lo, hi);
        if (j < k) {
            lo = j + 1;
        } else if (j > k) {
            hi = j - 1;
        } else {
            break;
        }
    }
    return nums[k];
}

private int partition(int[] a, int lo, int hi) {

    int i = lo;
    int j = hi + 1;
    while(true) {
        while(i < hi && less(a[++i], a[lo]));
        while(j > lo && less(a[lo], a[--j]));
        if(i >= j) {
            break;
        }
        exch(a, i, j);
    }
    exch(a, lo, j);
    return j;
}

private void exch(int[] a, int i, int j) {
    final int tmp = a[i];
    a[i] = a[j];
    a[j] = tmp;
}

private boolean less(int v, int w) {
    return v < w;
}