LeetCode #703 数据流中的第K大元素

#703 Kth Largest Element in a Stream

Posted by MatthewHan on 2020-05-30

Problem Description

设计一个找到数据流中第K大元素的类(class)。注意是排序后的第K大元素,不是第K个不同的元素。

你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器,它包含数据流中的初始元素。每次调用 KthLargest.add,返回当前数据流中第K大的元素。

note

你可以假设 nums 的长度≥ k-1k ≥ 1

e.g.

1
2
3
4
5
6
7
8
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

Solution

1. 额外数组空间

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class KthLargest {

private int len;
private final int index;
private final List<Integer> newData;

public KthLargest(int k, int[] nums) {
index = k;
len = nums.length;
newData = new ArrayList<>(nums.length);
for (int num : nums) {
newData.add(num);
}
}

public int add(int val) {
len++;
newData.add(val);
Collections.sort(newData);
return newData.get(len - index);
}
}

不过还是用原始数组,更底层更快一点:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class KthLargest {

private int len;
private int[] data;
private final int index;

public KthLargest(int k, int[] nums) {
index = k;
len = nums.length;
data = nums;
}

public int add(int val) {
len++;
int[] temp = new int[len];
System.arraycopy(data, 0, temp, 0, data.length);
temp[len - 1] = val;
Arrays.sort(temp);
data = new int[len];
System.arraycopy(temp, 0, data, 0, temp.length);
return data[len - index];
}
}

2. 小顶堆

题目中表述了add方法返回第K大元素即可,所以其实我们不必维护整个数组,只需要维护K个从大到小的元素即可。表面add方法而已,不是真的新增。

什么是堆?堆是一种非线性结构,可以把堆看作一个数组,也可以被看作一个完全二叉树,通俗来讲堆其实就是利用完全二叉树的结构来维护的一维数组

按照堆的特点可以把堆分为大顶堆小顶堆

  1. 大顶堆:每个结点的值都大于等于其左右孩子结点的值
  2. 小顶堆:每个结点的值都小于等于其左右孩子结点的值

大顶堆和小顶堆

直接说用法:

  • 升序:使用大顶堆

  • 降序:使用小顶堆

我们这边每次add方法需要返回第K个从大到小元素,所以每次执行add方法需要保留K个降序的元素,这里就是采用了小顶堆。我们可以用优先队列PriorityQueue<Integer>来做(大顶堆的话需要重写排序方法,默认是升序)。优先队列的作用是保证每次取出的元素都是队列中权值最小的。

当堆内元素的个数小于等于K时,调用add方法入队;当第K+1个元素想入队时,和队首比较大小(队首是最小的)。比队首大,队首元素出队,新元素入队;比队首小,不作任何操作。这样就能保证队列的元素个数是K个,并且是队首是最小的(不一定是从小到大排序,但是队首一定是整个队列里最小的)。

所以就是两步:

  1. 控制队列大小为K
  2. 保证队首的元素比新加入的元素大

所以整个数组的第K大元素就是队列的队首peek()方法。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class KthLargest {

// 小顶堆
private PriorityQueue<Integer> q;
private int k;

public KthLargestElementInaStream(int k, int[] nums) {
this.k = k;
q = new PriorityQueue<Integer>(k);
for (int i : nums) {
add(i);
}
}
public int add(int val) {
if (q.size() < k) {
q.offer(val);

} else if (q.peek() < val) {
q.poll();
q.offer(val);
}
return q.peek();
}
}

补充

PriorityQueue<Integer> queue = new PriorityQueue<>(1);这样一个优先队列,将数组[8, 5, 4, 12, 3]按顺序执行offer(element)方法,打印的结果是[3, 4, 5, 12, 8],在执行offer(1),打印的结果是[1, 4, 3, 12, 8, 5]。排序就是在二叉树的基础上,叶子节点和父节点比大小,比父节点大,交换位置,直到比父节点小或到根节点。