計算機程式的思維邏輯 (47) - 堆和PriorityQueue的應用

swiftma發表於2016-11-21

本系列文章經補充和完善,已修訂整理成書《Java程式設計的邏輯》(馬俊昌著),由機械工業出版社華章分社出版,於2018年1月上市熱銷,讀者好評如潮!各大網店和書店有售,歡迎購買:京東自營連結

計算機程式的思維邏輯 (47) -  堆和PriorityQueue的應用

45節介紹了堆的概念和演算法,上節介紹了Java中堆的實現類PriorityQueue,PriorityQueue除了用作優先順序佇列,還可以用來解決一些別的問題,45節提到了如下兩個應用:

  • 求前K個最大的元素,元素個數不確定,資料量可能很大,甚至源源不斷到來,但需要知道到目前為止的最大的前K個元素。這個問題的變體有:求前K個最小的元素,求第K個最大的,求第K個最小的。
  • 求中值元素,中值不是平均值,而是排序後中間那個元素的值,同樣,資料量可能很大,甚至源源不斷到來。

本節,我們就來探討如何解決這兩個問題。

求前K個最大的元素

基本思路

一個簡單的思路是排序,排序後取最大的K個就可以了,排序可以使用Arrays.sort()方法,效率為O(N*log2(N))。不過,如果K很小,比如是1,就是取最大值,對所有元素完全排序是毫無必要的。

另一個簡單的思路是選擇,迴圈選擇K次,每次從剩下的元素中選擇最大值,這個效率為O(N*K),如果K的值大於log2(N),這個就不如完全排序了。

不過,這兩個思路都假定所有元素都是已知的,而不是動態新增的。如果元素個數不確定,且源源不斷到來呢?

一個基本的思路是維護一個長度為K的陣列,最前面的K個元素就是目前最大的K個元素,以後每來一個新元素的時候,都先找陣列中的最小值,將新元素與最小值相比,如果小於最小值,則什麼都不用變,如果大於最小值,則將最小值替換為新元素。

這有點類似於生活中的末尾淘汰,新元素與原來最末尾的比即可,要麼不如最末尾,上不去,要麼替掉原來的末尾。

這樣,陣列中維護的永遠是最大的K個元素,而且不管源資料有多少,需要的記憶體開銷是固定的,就是長度為K的陣列。不過,每來一個元素,都需要找最小值,都需要進行K次比較,能不能減少比較次數呢?

解決方法是使用最小堆維護這K個元素,最小堆中,根即第一個元素永遠都是最小的,新來的元素與根比就可以了,如果小於根,則堆不需要變化,否則用新元素替換根,然後向下調整堆即可,調整的效率為O(log2(K)),這樣,總體的效率就是O(N*log2(K)),這個效率非常高,而且儲存成本也很低。

使用最小堆之後,第K個最大的元素也很容易獲得,它就是堆的根。

理解了思路,下面我們來看程式碼。

實現程式碼

我們來實現一個簡單的TopK類,程式碼如下所示:

public class TopK <E> {
    private PriorityQueue<E> p;
    private int k;
    
    public TopK(int k){
        this.k = k;
        this.p = new PriorityQueue<>(k);
    }

    public void addAll(Collection<? extends E> c){
        for(E e : c){
            add(e);
        }
    }
    
    public void add(E e) {
        if(p.size()<k){
            p.add(e);
            return;
        }
        
        Comparable<? super E> head = (Comparable<? super E>)p.peek();
        if(head.compareTo(e)>0){
            //小於TopK中的最小值,不用變
            return;
        }
        //新元素替換掉原來的最小值成為Top K之一。
        p.poll();
        p.add(e);
    }
    
    public <T> T[] toArray(T[] a){
        return p.toArray(a);
    }

    public E getKth(){
        return p.peek();
    }
}    
複製程式碼

我們稍微解釋一下。

TopK內部使用一個優先順序佇列和k,構造方法接受一個引數k,使用PriorityQueue的預設構造方法,假定元素實現了Comparable介面。

add方法,實現向其中動態新增元素,如果元素個數小於k直接新增,否則與最小值比較,只在大於最小值的情況下新增,新增前,先刪掉原來的最小值。addAll方法迴圈呼叫add方法。

toArray方法返回當前的最大的K個元素,getKth方法返回第K個最大的元素。

我們來看一下使用的例子:

TopK<Integer> top5 = new TopK<>(5);
top5.addAll(Arrays.asList(new Integer[]{
        100, 1, 2, 5, 6, 7, 34, 9, 3, 4, 5, 8, 23, 21, 90, 1, 0
}));

System.out.println(Arrays.toString(top5.toArray(new Integer[0])));
System.out.println(top5.getKth());
複製程式碼

保留5個最大的元素,輸出為:

[21, 23, 34, 100, 90]
21
複製程式碼

程式碼比較簡單,就不解釋了。

求中值

基本思路

中值就排序後中間那個元素的值,如果元素個數為奇數,中值是沒有歧義的,但如果是偶數,中值可能有不同的定義,可以為偏小的那個,也可以是偏大的那個,或者兩者的平均值,或者任意一個,這裡,我們假定任意一個都可以。

一個簡單的思路是排序,排序後取中間那個值就可以了,排序可以使用Arrays.sort()方法,效率為O(N*log2(N))。

不過,這要求所有元素都是已知的,而不是動態新增的。如果元素源源不斷到來,如何實時得到當前已經輸入的元素序列的中位數?

可以使用兩個堆,一個最大堆,一個最小堆,思路如下:

  1. 假設當前的中位數為m,最大堆維護的是<=m的元素,最小堆維護的是>=m的元素,但兩個堆都不包含m。
  2. 當新的元素到達時,比如為e,將e與m進行比較,若e<=m,則將其加入到最大堆中,否則將其加入到最小堆中。
  3. 第二步後,如果此時最小堆和最大堆的元素個數的差值>=2 ,則將m加入到元素個數少的堆中,然後從元素個數多的堆將根節點移除並賦值給m。

我們通過一個例子來解釋下,比如輸入元素依次為:

34, 90, 67, 45,1
複製程式碼

輸入第一個元素時,m即為34。

輸入第二個元素時,90大於34,加入最小堆,中值不變,如下所示:

計算機程式的思維邏輯 (47) -  堆和PriorityQueue的應用

輸入第三個元素時,67大於34,加入最小堆,但加入最小堆後,最小堆的元素個數為2,需調整中值和堆,現有中值34加入到最大堆中,最小堆的根67從最小堆中刪除並賦值給m,如下圖所示:

計算機程式的思維邏輯 (47) -  堆和PriorityQueue的應用

輸入第四個元素45時,45小於67,加入最大堆,中值不變,如下圖所示:

計算機程式的思維邏輯 (47) -  堆和PriorityQueue的應用

輸入第五個元素1時,1小於67,加入最大堆,此時需調整中值和堆,現有中值67加入到最小堆中,最大堆的根45從最大堆中刪除並賦值給m,如下圖所示:

計算機程式的思維邏輯 (47) -  堆和PriorityQueue的應用

實現程式碼

理解了基本思路,我們來實現一個簡單的中值類Median,程式碼如下所示:

public class Median <E> {
    private PriorityQueue<E> minP; // 最小堆
    private PriorityQueue<E> maxP; //最大堆
    private E m; //當前中值
    
    public Median(){
        this.minP = new PriorityQueue<>();
        this.maxP = new PriorityQueue<>(11, Collections.reverseOrder());
    }
    
    private int compare(E e, E m){
        Comparable<? super E> cmpr = (Comparable<? super E>)e;
        return cmpr.compareTo(m);
    }
    
    public void add(E e){
        if(m==null){ //第一個元素
            m = e;
            return;
        }
        if(compare(e, m)<=0){
            //小於中值, 加入最大堆
            maxP.add(e);
        }else{
            minP.add(e);
        }
        if(minP.size()-maxP.size()>=2){
            //最小堆元素個數多,即大於中值的數多
            //將m加入到最大堆中,然後將最小堆中的根移除賦給m
            maxP.add(this.m);
            this.m = minP.poll();
        }else if(maxP.size()-minP.size()>=2){
            minP.add(this.m);
            this.m = maxP.poll();
        }
    }
    
    public void addAll(Collection<? extends E> c){
        for(E e : c){
            add(e);
        }
    }
    
    public E getM() {
        return m;
    }
}
複製程式碼

程式碼和思路基本是對應的,比較簡單,就不解釋了。我們來看一個使用的例子:

Median<Integer> median = new Median<>();
List<Integer> list = Arrays.asList(new Integer[]{
        34, 90, 67, 45, 1, 4, 5, 6, 7, 9, 10
});
median.addAll(list);
System.out.println(median.getM());
複製程式碼

輸出為中值9。

小結

本節介紹了堆和PriorityQueue的兩個應用,求前K個最大的元素和求中值,介紹了基本思路和實現程式碼,相比使用排序,使用堆不僅實現效率更高,而且還可以應對資料量不確定且源源不斷到來的情況,可以給出實時結果。

到目前為止,我們介紹了佇列的兩個實現,LinkedList和PriortiyQueue,Java容器類中還有一個佇列的實現類ArrayDeque,它是基於陣列實現的,我們知道,一般而言,因為需要移動元素,陣列的插入和刪除效率比較低,但ArrayDeque的效率卻很高,甚至高於LinkedList,它是怎麼實現的呢?讓我們下節來探討。


未完待續,檢視最新文章,敬請關注微信公眾號“老馬說程式設計”(掃描下方二維碼),深入淺出,老馬和你一起探索Java程式設計及計算機技術的本質。用心原創,保留所有版權。

計算機程式的思維邏輯 (47) -  堆和PriorityQueue的應用

相關文章