從排序算法到TopK問題

一、前言

排序算法大家都很熟悉了，解法多種多樣。

有一個問題和排序算法很相近，TopK問題：從N個數中選出最大的K個數，N通常遠大于K。

總結了一些解法，供大家參考。

二、冒泡

private static float[] pickTopKByBubbleSort(float[] a, int k) {
        int n = a.length;
        for (int i = 0; i < k; i++) {
            for (int j = 0; j < n - i - 1; j++) {
                if (a[j] > a[j + 1]) {
                    float t = a[j];
                    a[j] = a[j + 1];
                    a[j + 1] = t;
                }
            }
        }

        float[] r = new float[k];
        for (int i = 0; i < k; i++) {
            r[i] = a[n - i - 1];
        }
        return r;
    }           

冒泡排序的要義在于：外循環為排序趟數，内循環為每趟比較的次數，每趟得到一個最大的數，放到這一趟的末端。

如果是全數組排序的話，複雜度為O(n^2), 如果隻選TopK的話，複雜度為O(n*k)。

三、堆

private static float[] pickTopKByHeapSort(float a[], int k) {
        int n = a.length;
        PriorityQueue<Float> queue = new PriorityQueue<>(n);
        for (int i = 0; i < k; i++) {
            queue.offer(a[i]);
        }

        for (int i = k; i < n; i++) {
            if (a[i] > queue.peek()) {
                queue.poll();
                queue.offer(a[i]);
            }
        }

        float[] r = new float[k];
        for (int i = k - 1; i >= 0; i--) {
            r[i] = queue.poll();
        }
        return r;
    }           

堆的實作，在JDK中對應的是PriorityQueue，預設情況下是最小堆。

元素插入和删除的時間複雜度都是O(logn）。

堆可以用于排序：将所有元素插入堆中，在依次取出，即可得到有序排列的元素，時間複雜度為O(nlogn)。

如果按照上面冒泡排序的做法，将所有元素插入，再取出K個：

插入時間複雜度為O(nlogn), 取出的時間複雜度為O(klogn), 總體時間複雜度還是O(nlogn)。

當然，雖然複雜度都是O(nlogn)，但是比全體元素排序快一點（不用取盡n個元素）。

堆在TopK問題上更優的解法是：

1、隻放k個元素到堆中，然後周遊剩下的元素，和堆頂的元素（堆中最小的值）比較，

2、若大于堆頂，則取出堆頂的元素，插入目前元素；若小于等于堆頂，直接跳過；

3、最後依次取出堆中元素。

平均時間複雜度為O(nlogk);

最快的情況下是O(n)，如果數組本身已經有序且是逆序的話。

四、快排

先來看下快排的實作：

private static int partition(float[] a, int left, int right) {
        int i = left;
        int j = right;
        float key = a[left];

        while (i < j) {
            while (i < j && a[j] <= key) {
                j--;
            }
            a[i] = a[j];
            while (i < j && a[i] >= key) {
                i++;
            }
            a[j] = a[i];
        }
        a[i] = key;
        return i;
    }

    private static void sort(float[] a, int left, int right) {
        if (left >= right) {
            return;
        }
        int index = partition(a, left, right);
        sort(a, left, index - 1);
        sort(a, index + 1, right);
    }           

快排的基本思想是：

1、通過一趟排序将要資料分割成獨立的兩部分，其中一部分的所有資料都比另外一部分的資料都要小；

2、然後再按此方法對這兩部分資料分别進行第1步的操作（遞歸），以此達到整個資料變成有序序列。

在TopK問題上，可以借鑒以上第1點，以及第2點的一部分：遞歸。

不同之處在于，求解TopK不需要對整個資料排序，隻需要最大的K個數移動到數組的第K個位置的前面即可。

代碼如下：

private static void sort(float[] a, int left, int right, int k) {
        if (left >= right) {
            return;
        }
        int index = partition(a, left, right);
        int len = index - left + 1;
        if (k < len) {
            sort(a, left, index - 1, k);
        } else if (k > len) {
            sort(a, index + 1, right, k - len);
        }
    }

    private static float[] pickTopKByQuickSort(float a[], int k) {
        sort(a, 0, a.length - 1, k);
        float[] r = new float[k];
        for (int i = 0; i < k; i++) {
            r[i] = a[i];
        }
        // 執行以上操作之後，r[]數組确實包含了top k的元素，但是不一定有序
        // 如果僅僅是要求TopK而不需要有序，則不需要執行下面的排序函數
        sort(r, 0, k - 1);
        return r;
    }           

最壞的情況下，時間複雜度和快排一樣: O(n^2)。

平均情況下，快排的時間複雜度為O(nlogn)，求解TopK時間複雜度為O(n)。

五、桶排序

冒泡排序，堆排序，快速排序等都是基于比較的排序，時間複雜度下限為O(nlogn)。

而有一些排序的時間複雜度可以是O(n), 比如桶排序和基數排序。

但是這類排序有較為苛刻的限制條件：元素須是數值類型。

同時，如果是桶排序的話，範圍不能太大，否則需要的桶太多，空間複雜度無法接受。

桶排序的元素不一定要是整數，有時候浮點數也可以。

比如，如果資料取值範圍在[0, 100.0], 并且隻有兩位小數, 可以這麼解：

private static float[] pickTopKByBucketSort(float[] a, int k) {
        int[] bucket = new int[10000];

        for (int i = 0; i < a.length; i++) {
            int index = Math.round(a[i] * 100f);
            bucket[index]++;
        }

        float[] r = new float[k];
        int p = 0;
        for (int i = bucket.length - 1; i >= 0 && p < k; i--) {
            int c = bucket[i];
            while (c > 0 && p < k) {
                r[p++] = i / 100f;
                c--;
            }
        }

        return r;
    }           

取一萬個桶，将每個元素乘以100，四舍五入為整數（有一些小數在計算機二進制中無法精确表示，同時直接強轉成int的話是向下取整），

放到對應的桶中；

然後從最大的桶開始檢索，收集100個元素。

時間複雜度跟資料量和桶數量（取決于資料的大小範圍）有關。

在資料量遠大于桶數量時，時間複雜度為O(n)。

六、測試

內建測試的代碼如下：

public static void main(String[] args) {
        int n = 100000;
        int k = 100;
        boolean success = true;
        
        Random r = new Random();
        for (int j = 0; j < 100; j++) {
            float[] a = new float[n];
            for (int i = 0; i < n; i++) {
                float t = r.nextFloat() * 100f;
                t = (int) (t * 100) / 100f;
                a[i] = t;
            }
            
            float[] tmp = Arrays.copyOf(a, n);
            float[] top1 = pickTopKByBubbleSort(tmp, k);
   
            tmp = Arrays.copyOf(a, n);
            float[] top2 = pickTopKByHeapSort(tmp, k);
       
            tmp = Arrays.copyOf(a, n);
            float[] top3 = pickTopKByQuickSort(tmp, k);
            
            tmp = Arrays.copyOf(a, n);
            float[] top4 = pickTopKByBucketSort(tmp, k);

            if (!Arrays.equals(top1, top2)
                    || !Arrays.equals(top1, top3)
                    || !Arrays.equals(top1, top4)) {
                success = false;
                break;
            }
        }

        if (success) {
            System.out.println("success");
        } else {
            System.out.println("failed");
        }
    }           

七、總結

這麼多種方案，那種方案比較好呢？這個要看情況。

• 冒泡的方法看起來時間複雜度是最高的，但是假如K很小，比如極端情況 , K=1，這時候冒泡的方法反而是最快的；
• 平均情況下，如果可以用的上桶排序（資料取值範圍較小），那桶排序是最快的；

• 如果數值取值範圍較小，而N很大（比方說100000），K又比較小（比方說100）時，運作發現堆的方法比快排的方法快。

  原因猜測：範圍小，資料多，應該有很多數值是重複的； 用堆的方法，堆會很快攢到Top的數值（即使還不是TopK），檢索到那些小的數值的時候就不需要插入，進而檢索後面的元素時，操作就隻剩簡單的周遊了。

跳脫排序算法和TopK問題本身，最主要的還是思想。

例如，日常項目中，堆（優先隊列）就不時會遇到，這前提是腦海中知道有優先隊列這個東西，知道它的特性。

有的知識就是這樣的，專門去學的時候不知道有什麼用，但或許某一天碰到了，就會想起它了。