1. 程式人生 > >Top K以及java priorityqueue

Top K以及java priorityqueue

vat 元素 poll() 封裝 輸入 現在 pre poll 判斷

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     }

Top K以及java priorityqueue