如何解決TOP-K問題

前言:最近在開發一個功能:動態展示的訂單數量排名前10的城市,這是一個典型的Top-k問題,其中k=10,也就是說找到一個集合中的前10名。實際生活中Top-K的問題非常廣泛,比如:微博熱搜的前100名、抖音直播的小時榜前50名、百度熱搜的前10條、博客園點贊最多的blog前10名,等等如何解決這類問題呢?初步的想法是將這個數據集合排序,然後直接取前K個返回。這樣解法可以,但是會存在一個問題:排序了很多不需要去排序的數據,時間複雜度過高.假設有數據100萬,對這個集合進行排序需要很長的時間,即便使用快速排序,時間複雜度也是O(nlogn),那麼這個問題如何解決呢?解決方法就是以空間換時間,使用優先級隊列

目錄

一:  認識PriorityQueue
二:利用PriorityQueue解決topk問題
三:總結

 

一:認識PriorityQueue

1.1:PriorityQueue位於java.util包下,繼承自Collection,因此它具有集合的屬性,並且繼承自Queue隊列,擁有add/offer/poll/peek等一系列操作元素的能力,它的默認大小是11,底層使用Object[] 來保存元素,數組的話肯定會有擴容,當添加元素的時候大小超過數組的容量,就會擴容,擴容的大小為原數組的大小加上,如果元素的數量小於64,則每次加2,如果大於64,則每次增加一半的容量。

 1.2:PriorityQueue的構造方法

    public PriorityQueue(Comparator<? super E> comparator) {
        this(DEFAULT_INITIAL_CAPACITY, comparator);
    }
 public PriorityQueue(int initialCapacity,
                         Comparator<? super E> comparator) {
        // Note: This restriction of at least one is not actually needed,
        // but continues for 1.5 compatibility
        if (initialCapacity < 1)
            throw new IllegalArgumentException();
        this.queue = new Object[initialCapacity];
        this.comparator = comparator;
    }

比較常用的就是這兩個構造方法,其中第一個構造方法中需要構造一個比較器,第二個構造方法添加初始容量和比較器,比較器可以自定義任何元素的優先級,按照需要增加元素的優先級展示

 1.3:PriorityQueue的常用API

1.3.1:offer方法和add方法用於添加元素,本質上offer方法和add方法是相同的:

    public boolean add(E e) {
        return offer(e);
    }

 offer方法主要步驟就是判空、擴容、添加元素,添加元素的話,siftup方法里會根據構造方法,如果有比較器就進行比較,沒有比較器的話就給元素賦予比較能力,並且根據構造的大小,也就是

 initialCapacity進行比較,如果比較器的compare方法不符合定義的規則,直接break;符合的話會給數組的元素進行賦值

public boolean offer(E e) {
        if (e == null)
            throw new NullPointerException();
        modCount++;
        int i = size;
        if (i >= queue.length)
            grow(i + 1);
        size = i + 1;
        if (i == 0)
            queue[0] = e;
        else
            siftUp(i, e);
        return true;
    }

1.3.2:poll方法和peek方法都是返回頭元素,不同之處在於poll方法會返回頭頂元素並且移除元素,peek方法不會移除頭頂元素:

  public E poll() {
        if (size == 0)
            return null;
        int s = --size;
        modCount++;
        E result = (E) queue[0];
        E x = (E) queue[s];
        queue[s] = null;
        if (s != 0)
            siftDown(0, x);
        return result;
    }
    public E peek() {
        return (size == 0) ? null : (E) queue[0];
    }

 二:PriorityQueue解決問題

2.1:數組的前K大值

代碼:

import java.util.PriorityQueue;

public class TopK {

    //找出前k個最大數,採用小頂堆實現
    public static int[] findKMax(int[] nums, int k) {

        PriorityQueue<Integer> pq = new PriorityQueue<>(k);//隊列默認自然順序排列,小頂堆,不必重寫compare

        for (int num : nums) {
            if (pq.size() < k) {
                pq.offer(num);
                //如果堆頂元素 < 新數,則刪除堆頂,加入新數入堆,保持堆中存儲着大值
            } else if (pq.peek() < num) {
                pq.poll();
                pq.offer(num);
            }
        }

        int[] result = new int[k];
        for (int i = 0; i < k && !pq.isEmpty(); i++) {
            result[i] = pq.poll();
        }
        return result;
    }
}

 測試:

 public static void main(String[] args) {
        int[] arr = new int[]{1, 6, 2, 3, 5, 4, 8, 7, 9};
        System.out.println(Arrays.toString(findKMax(arr, 3)));
    }

//輸出:

優先級隊列是如何解決這個問題的呢?

PriorityQueue默認是小頂堆,那麼什麼是小頂堆?什麼是大頂堆?假設有3、6、8三個數,需要存儲在優先級隊列里,畫個圖大家理解下:

 可以看出小頂堆的頭頂元素存儲着整個數據集合中數字最小的元素,而大頂堆存儲着整個數據集合中數字最大的元素,也就是一個按照升序排列,一個按照降序排列:

//小頂堆的構建方法:
PriorityQueue<Integer> queue = new PriorityQueue<>(k);
//這種寫法等價於: 
PriorityQueue<Integer> queue = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o1-o2;
            }
        });
//同時等價於(lamda表達式的寫法)
PriorityQueue<Integer> queue = new PriorityQueue<>(k, (o1, o2) -> o1-o2);

 //大頂堆的構建方法:
 PriorityQueue<Integer> queue = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2-o1;
            }
        });

拿測試用例這個例子來說:

構建的是指定容量的小頂堆,因此每次queue.peek()返回的是最小的數字,在遍曆數組的過程中,如果遇到比該數字大的元素就將最小的數字poll(移除掉),然後將較大的元素添加到堆中,在添加進去堆中的時候,堆同時會按照優先級比較,將最小的元素再次放到堆頂,這樣的做法就是會一直保持堆中的元素是相對較大的,同時堆頂元素是堆中最小的。

按照測試用例給出的例子,{1, 6, 2, 3, 5, 4, 8, 7, 9} 優先級隊列將會是這樣轉變的:(注意:本質上優先級隊列的實現方式是數組,這裡只是用二叉樹的方式表現出來)

 

 假如該題換個角度,求出現頻率最低的元素怎麼做呢?

相信你根據上面的講述應該也明白了:直接構建一個大頂堆,這樣元素最大的值在堆頂,每次去和數組的元素的值去做比較,只要堆頂元素比數組的值小,就將堆頂元素poll出來,然後將數組的值添加進去,這樣就可以一直保持集合數組中一直是最小的k個數字。

 2.2:前k個高頻元素

當k = 1 時問題很簡單,線性時間內就可以解決,只需要用哈希表維護元素出現頻率,每一步更新最高頻元素即可。當 k > 1 就需要一個能夠根據出現頻率快速獲取元素的數據結構,這裡就需要用到優先隊列

首先建立一個元素值對應出現頻率的哈希表,使用 HashMap來統計,然後構建優先級隊列,這裡依舊是構建小頂堆,不過因為該題是計算元素出現的頻率,因此我們需要將每個元素的頻率值做對比,

需要重寫優先級隊列的comparator,需要手工填值:這個步驟需要 O(N)O(N) 時間其中 NN 是列表中元素個數。

第二步建立堆,堆中添加一個元素的複雜度是 O(\log(k))O(log(k)),要進行 NN 次複雜度是 O(N)O(N)。

最後一步是輸出結果,複雜度為 O(k\log(k))O(klog(k)):

public int[] topK(int[] nums, int k) {
        //統計字符出現的頻率的map
        Map<Integer, Integer> count = new HashMap();
        for (int num : nums) {
            count.put(num, count.getOrDefault(num, 0) + 1);
        }

        //根據出現頻率的map來構建k個元素的優先級隊列
        PriorityQueue<Integer> heap =
                new PriorityQueue<>(k, (o1, o2) -> count.get(o1) - count.get(o2));

        for (int n : count.keySet()) {
            heap.add(n);
            if (heap.size() > k)
                heap.poll();
        }

        int[] result = new int[heap.size()];
        for (int i = 0; i < result.length; i++) {
            result[i] = heap.poll();
        }
        return result;

    }

 測試:

public static void main(String[] args) {
        int[] nums = {1, 1, 1, 3, 5, 6, 5, 6, 5, 4, 7, 8};
        int[] res = new TOPK().topK(nums, 3);
        System.out.println(Arrays.toString(res));
    }

//輸出

 

  對於這樣的問題需要先對原數組進行處理,比如在計算前3頻率這個問題上,我們需要先計算數組中數字出現的頻率,然後維護一個哈希表用來存儲元素的頻率。對於類似問題:微博熱搜的前10名,那麼肯定需要統計搜索頻次,抖音小時榜前10名,那麼肯定要統計要計算時段的觀看人數,優先隊列只不過是一個存儲K元素的一個容器,它不負責統計,只負責維護一個K元素的最大或者最小堆頂,對於數據採用什麼樣的優先級順序需要自定義。

三:總結

      本篇博客主要介紹了我們在實際中遇見的TOP-K問題有哪些,以及優先級隊列PriorityQueue的基本原理介紹,接着由易到難的講解了如何通過優先級隊列PriorityQueue來解決TOP-k問題,這兩個問題都比較經典。對於理解優先級隊列的含義、以及為什麼它能解決該問題,想明白這點很重要。希望大家能夠做到舉一反三,下次面對同等問題的時候,能順序解決。起碼棘手的topk問題對於我們來說,有個PriorityQueue這個神兵利器,就顯得很簡單easy咯

最後如果對學習java有興趣可以加入群:618626589,本群旨在打造無培訓廣告、無閑聊扯皮、無注水斗圖的純技術交流群,群里每天會分享有價值的問題和學習資料,歡迎各位隨時加入~

Tags: