[每日一題] 第二十題：最小的k個數

輸入整數陣列 arr ，找出其中最小的 k 個數。例如，輸入4、5、1、6、2、7、3、8這8個數字，則最小的4個數字是1、2、3、4。

示例 1：

輸入：arr = [3,2,1], k = 2
輸出：[1,2] 或者 [2,1]

示例 2：

輸入：arr = [0,1,2,1], k = 1
輸出：[0]

限制：

• 0 <= k <= arr.length <= 10000
• 0 <= arr[i] <= 10000

方法一：使用 Arrays.sort() 方法

程式碼

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
       Arrays.sort(arr);
        int[] result = new int[k];
        for (int i = 0; i < k ; i++) {
            result[i] = arr[i];
        }

        return result;
    }
}

時間複雜度

方法二：堆排序

比較直觀的想法是使用堆資料結構來輔助得到最小的 K 個數，堆的性質是每次可以得到最大或最小的元素。我們可以使用一個大小為 K 的最大堆（大頂堆），將陣列中的元素依次推入堆，當堆的大小超過 K 時，便將多餘的元素從堆頂彈出。我們以陣列 [5, 4, 1, 3, 6, 2, 9]， k=3 為例展示元素入堆的過程，如下面動圖所示：

這樣，由於每次從堆頂彈出的數都是堆中最大的，最小的 K 個元素一定會留在堆裡。 這樣，把陣列中的元素全部入堆之後，堆中剩下的 K 個元素就是最小的 K 個數了。

注意在動畫中，我們並沒有畫出堆的內部結構，因為這部分內容並不重要，我們只需要知道堆每次會彈出最大的元素即可。在寫程式碼的時候，我們使用的也是庫函式中優先佇列資料結構，如 Java 中的 priorityQueue。在面試中，我們不需要實現堆的內部結構，把資料結構使用好，會分析複雜度即可。

以下是題解程式碼。感謝評論區提醒，這裡的程式碼可以做一些最佳化，如果當前數字不小於堆頂元素，數字可以直接丟掉，不如堆，下方的程式碼已更新：

public int[] getLeastNumbers(int[] arr, int k) {
    if (k == 0) {
        return new int[0];
    }
    // 使用一個最大堆（大頂堆）
    // Java 的 PriorityQueue 預設是小頂堆，新增 comparator 引數使其變成最大堆
    Queue<Integer> heap = new PriorityQueue<>(k, (i1, i2) -> Integer.compare(i2, i1));

    for (int e : arr) {
        // 當前數字小於堆頂元素才會入堆
        if (heap.isEmpty() || heap.size() < k || e < heap.peek()) {
            heap.offer(e);
        }
        if (heap.size() > k) {
            heap.poll(); // 刪除堆頂最大元素
        }
    }

    // 將堆中的元素存入陣列
    int[] res = new int[heap.size()];
    int j = 0;
    for (int e : heap) {
        res[j++] = e;
    }
    return res;
}

複雜度

• 空間複雜度：由於使用了一個大小為 K 的堆，空間複雜度為 O(N)
• 時間複雜度：入堆和出堆的時間複雜度均為 O(logK)，每個元素都需要進行一次入堆操作，故演算法的時間複雜度為 O(n logK)。

來源

作者：nettee
連結：leetcode-cn.com/problems/zui-xiao-...
來源：力扣（LeetCode）

方法三：快排變形

Top K 問題的另一個解法就比較難想到，需要在平時有演算法的積累，實際上，”查詢第 K 大元素“是一類演算法問題，稱為選擇問題。找到第 K 大的元素，或者是前 K 大的元素，有一個經典的 quick select（快速選擇）演算法。這個名字和 quick sort（快速排序）看起來很像，演算法的思想野核快速排序相似，都是分治法的思想。

這個 partition 操作是原地進行的，需要 O(N) 的時間，接下來，快速排序會遞迴的排序左右兩側的資料，而 quick select（快速選擇）的演算法不同之處在於，接下來只需要遞迴的選擇一側的陣列，快速選擇演算法相當於一個”不完全“的快速排序，因為我們只需要知道最小的 k 個數是哪些，並不需要知道他們的順序。

我們的目的是尋找最小的 k 個數，假設經過一次 partition 操作，樞紐元素位於下標 m，也就是說，左側的陣列有 m 個元素，是原陣列中的最小的 m 個數。那麼：

• 若 k=m，我們就找到了最小的 k 個數，就是左側的資料。
• 若 k<m，則最小的 k 個數一定都在左側陣列中，我們只需要對左側陣列遞迴地 partition 即可。
• 若 k>m，則左側陣列的 m 個數都屬於最小的 k 個數，我們還需要在右側陣列中尋找最小的 k-m 個數，對右側陣列遞迴的 partition 即可。

這種方法需要多加領會思想，如果你對快速排序掌握得很好，那麼稍加推導應該不難掌握 quick select 的要領。

以下是題解程式碼：

程式碼

public int[] getLeastNumbers(int[] arr, int k) {
    if (k == 0) {
        return new int[0];
    } else if (arr.length <= k) {
        return arr;
    }

    // 原地不斷劃分陣列
    partitionArray(arr, 0, arr.length - 1, k);

    // 陣列的前 k 個數此時就是最小的 k 個數，將其存入結果
    int[] res = new int[k];
    for (int i = 0; i < k; i++) {
        res[i] = arr[i];
    }
    return res;
}

void partitionArray(int[] arr, int lo, int hi, int k) {
    // 做一次 partition 操作
    int m = partition(arr, lo, hi);
    // 此時陣列前 m 個數，就是最小的 m 個數
    if (k == m) {
        // 正好找到最小的 k(m) 個數
        return;
    } else if (k < m) {
        // 最小的 k 個數一定在前 m 個數中，遞迴劃分
        partitionArray(arr, lo, m-1, k);
    } else {
        // 在右側陣列中尋找最小的 k-m 個數
        partitionArray(arr, m+1, hi, k);
    }
}

// partition 函式和快速排序中相同，具體可參考快速排序相關的資料
// 程式碼參考 Sedgewick 的《演算法4》
int partition(int[] a, int lo, int hi) {
    int i = lo;
    int j = hi + 1;
    int v = a[lo];
    while (true) { 
        while (a[++i] < v) {
            if (i == hi) {
                break;
            }
        }
        while (a[--j] > v) {
            if (j == lo) {
                break;
            }
        }

        if (i >= j) {
            break;
        }
        swap(a, i, j);
    }
    swap(a, lo, j);

    // a[lo .. j-1] <= a[j] <= a[j+1 .. hi]
    return j;
}

void swap(int[] a, int i, int j) {
    int temp = a[i];
    a[i] = a[j];
    a[j] = temp;
}

上述程式碼中需要注意一個細節（評論區有好幾個小夥伴問到，這裡補充說明一下）：

partitionArray 函式中，兩次遞迴呼叫傳入的引數為什麼都是 k？特別是第二個呼叫，我們在右側陣列中尋找最小的 k-m 個數，但是對於整個陣列而言，這是最小的 k 個數。所以說，函式呼叫傳入的引數應該為 k。

該程式碼的成績還是非常好的：

複雜度

• 空間複雜度 O(1)，不需要額外空間。
• 時間複雜度的分析方法和快速排序類似。由於快速選擇只需要遞迴一邊的陣列，時間複雜度小於快速排序，期望時間複雜度為 O(n)，最壞情況下的時間複雜度為 O(n^2)。

兩種方法的優劣性比較

在面試中，另一個常常問的問題就是這兩種方法有何優劣。看起來分治法的快速選擇演算法的時間、空間複雜度都優於使用堆的方法，但是要注意到快速選擇演算法的幾點侷限性：

第一，演算法需要修改原陣列，如果原陣列不能修改的話，還需要複製一份陣列，空間複雜度就上去了。

第二，演算法需要儲存所有的資料。如果把資料看成輸入流的話，使用堆的方法是來一個處理一個，不需要儲存資料，只需要儲存 k 個元素的最大堆。而快速選擇的方法需要先儲存下來所有的資料，再執行演算法。當資料量非常大的時候，甚至記憶體都放不下的時候，就麻煩了。所以當資料量大的時候還是用基於堆的方法比較好。

來源

作者：nettee
連結：leetcode-cn.com/problems/zui-xiao-...
來源：力扣（LeetCode）

來源：力扣（LeetCode）
連結：leetcode-cn.com/problems/zui-xiao-...

本作品採用《CC 協議》，轉載必須註明作者和本文連結