Top K以及java priorityqueue

標籤：code   實現   一個   parent   priority   思路   最大   integer   個數   

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             @Override3             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