Top K以及java priorityqueue

Posted 2020-12-26 Lin.B

Top K问题比较常见啦，这里总结一下方法。

1、用最小堆来做。

思路是先利用数组中前k个数字建一个最小堆，然后将剩余元素与堆顶元素进行比较，如果某个元素比堆顶元素大，就替换掉堆顶元素，并且重新调整成最小堆。

到这里，堆中保存着的其实是前k个最大的数字。堆顶就是第K个最大的数字。这样前k个，第k个都可以求出来了。代码如下：

 1     public void find(int[] nums, int k){
 2         PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();
 3         for ( int i = 0 ; i < k ; i ++ ) priorityQueue.offer(nums[i]);
 4         for ( int i = k ; i < nums.length ; i ++ ){
 5             if ( nums[i] > priorityQueue.peek() ){
 6                 priorityQueue.poll();
 7                 priorityQueue.offer(nums[i]);
 8             }
 9         }
10         System.out.println(priorityQueue.peek());
11         while ( priorityQueue.peek() != null ){
12             System.out.print(priorityQueue.poll()+"   ");
13         }
14     }

输出第一行是第k大的数，第二行是前k大的数。

现在有两个思考：1、这个代码是对重复数字有效的，也就是如果输入是[3,4,5,6,6]，求第三大的数字，结果就是5。但是实际上我们是想忽略重复，希望输出4，应该怎么办？2、如果是想求最小的k个数字应该怎么办。

如果想忽略重复，我们在建堆和后面元素入堆的过程中增加一个判断就可以了。

如果想要求前k小或者是d第k小的数字，这里就要了解一下java对堆的封装类：priorityqueue。这个封装类默认实现最小堆并且考虑重复。如果想实现上面的思路，就要重写compare方法。compare方法在siftup的过程中使用到，siftup这个函数就是向上调整，这个函数会比较孩子节点和父亲节点的大小关系是否满足要求，如果满足，就break；如果不满足就会交换孩子和父亲的值。

 1 private void siftUp(int k, E x) {
 2     if (comparator != null)
 3         siftUpUsingComparator(k, x);
 4     else
 5         siftUpComparable(k, x);
 6 }
 7 private void siftUpUsingComparator(int k, E x) {
 8     while (k > 0) {
 9         int parent = (k - 1) >>> 1;
10         Object e = queue[parent];
11         if (comparator.compare(x, (E) e) >= 0)
12             break;
13         queue[k] = e;
14         k = parent;
15     }
16     queue[k] = x;
17 }

注意这里如果没有重写compare方法，就会调用默认的siftupComparable方法（这个方法就是构建最小堆的方法）。如果重写了，那么就会使用compare方法中的比较方法来判断是否满足条件。我重写了一个最大堆的代码如下：

1 PriorityQueue<Integer> priorityQueue = new PriorityQueue<>(new Comparator<Integer>() {
2             @Override
3             public int compare(Integer o1, Integer o2) {
4                 return o2-o1;
5             }
6         });

compare方法中o1参数是孩子节点，o2参数是父亲节点。这里最大堆的意思就是：如果父亲节点的值大于孩子节点的值，就不需要调整；否则交换。所以根据这个原则建立起的堆就是最大堆。

那么求第k个最小的数，以及前k个最小的数代码就显而易见了：

 1     public void find(int[] nums, int k){
 2         PriorityQueue<Integer> priorityQueue = new PriorityQueue<>(new Comparator<Integer>() {
 3             @Override
 4             public int compare(Integer o1, Integer o2) {
 5                 return o2-o1;
 6             }
 7         });
 8         for ( int i = 0 ; i < k ; i ++ ) 
 9         {
10             if ( !priorityQueue.contains(nums[i]) ) priorityQueue.offer(nums[i]);
11         }
12         for ( int i = k ; i < nums.length  ; i ++ ){
13             if ( nums[i] < priorityQueue.peek() && !priorityQueue.contains(nums[i]) ){
14                 priorityQueue.poll();
15                 priorityQueue.offer(nums[i]);
16             }
17         }
18         System.out.println(priorityQueue.peek());
19         while ( priorityQueue.peek() != null ){
20             System.out.print(priorityQueue.poll()+"   ");
21         }
22     }