Python|算法|快速排序|如何在O(n)查找第K大元素

王爭老師講過，學習算法不是死記硬背一些源代碼或概念，而是學習算法的實現思路、思維、應用場景，從而達到靈活運用。

比如現在要時間復雜度為 O(n)，在一個長度為 n 的數組中查找到第 K 大的元素，你會怎么做呢？ 你可能會說這很簡單啊，第一次遍歷數組找到第 1 大元素，第二次遍歷找到第 2 大，...，第 K 次就可以找到第 K 大 但是這樣的時間復雜度就不是 O(n), 而是 K*O(n), 當 K 接近 n 時，時間復雜度就是 O(n^2)。 如果你運用快速排序算法的思想，你就可以在 O(n) 的時間復雜度內找到第 K 大元素。

快速排序算法

快速排序算法和歸並排序算法一樣，都是利用分治算法。但是它們卻有本質的不同，歸並排序是自下而上，先求解下面的子問題求，然后再逐層歸並，最后全部有序；而快速排序是自上而下，下面的子問題解決后，數據就全部有序。

快速排序的思路是這樣的，在數組中隨機選取一個數據，例如選取最后一個元素 m 做為分區元素，比 m 小的放 m 的左邊，反之放右邊，再分別對左右邊的分區再分別進行分區，直到分區元素縮小到 1 個，此時數據已經全部有序。
下面是我根據理解編寫的快速排序代碼（python 語言）

import random

def quick_sort(data_list):
    length = len(data_list)
    quick_sort_c(data_list,0,length-1)

def quick_sort_c(data_list,begin,end):
    """
    可以遞歸的函數調用
    """
    if begin >= end:
        return
    else:
        
        index = partition(data_list,begin,end)
        print(data_list)
        quick_sort_c(data_list,begin,index-1)
        quick_sort_c(data_list,index+1,end)

def partition(data_list,begin,end):
    
    partition_key = data_list[end]

    
    
    index = begin 
    for i in range(begin,end):
        if data_list[i] < partition_key: 
            data_list[i],data_list[index] = data_list[index],data_list[i] 
            index+=1

    data_list[index],data_list[end] = data_list[end],data_list[index] 
    return index

if __name__ == "__main__":
    data_list = [random.randint(0,100) for i in range(10)]
    
    print("原始數組:", data_list)
    print("排序過程如下")
    quick_sort(data_list)
    print("最終結果:",data_list)

執行結果如下所示：

原始數組: [66, 1, 10, 95, 87, 16, 14, 88, 87, 82]
排序過程如下
[66, 1, 10, 16, 14, 82, 87, 88, 87, 95]
[1, 10, 14, 16, 66, 82, 87, 88, 87, 95]
[1, 10, 14, 16, 66, 82, 87, 88, 87, 95]
[1, 10, 14, 16, 66, 82, 87, 88, 87, 95]
[1, 10, 14, 16, 66, 82, 87, 88, 87, 95]
[1, 10, 14, 16, 66, 82, 87, 88, 87, 95]
[1, 10, 14, 16, 66, 82, 87, 88, 87, 95]
[1, 10, 14, 16, 66, 82, 87, 88, 87, 95]
[1, 10, 14, 16, 66, 82, 87, 87, 88, 95]
[1, 10, 14, 16, 66, 82, 87, 87, 88, 95]
最終結果: [1, 10, 14, 16, 66, 82, 87, 87, 88, 95]

性能分析

快速排序是一種原地排序算法，不需要借助額外的存儲空間；由於分區的過程中由於其他元素的影響，在交換位置時會破壞原有的先后順序，比如 3，5，6，3，2 在第一次分區后，兩個 3 的相對次序已經改變，因此快速排序是一種不穩定的排序算法；時間復雜度為 O(nlogn)，但在極端情況下會降低到 O(n^2)，比如在數據已經是有序的情況時，需要進行 n 次分區，每次分區需要平均掃描 n/2 個元素，因此這種情況下時間復雜度為 O(n^2)。

O(n) 的時間內查找第 K 大元素的方法

通過觀察運行上面快速排序的過程可以發現，第一個分區鍵為 82，在第一次分區后，它是數組中的第 6 個元素，那么可以斷定，82 就是第 6 小元素，或者 82 就是第 （10-6+1）=5 大元素，需要查找最 3 大元素，那么這個元素一定在第一次分區的右部分進行分區操作，求得分區鍵的下標 index = n - K = 10 -3 = 7 時返回分區鍵即是所求得的數據。下面我通過代碼實現如下：

def find_top_k(data_list,K):
    length = len(data_list)
    begin = 0
    end = length-1
    index = partition(data_list,begin,end)
    while index != length - K:
        if index >length - K:
            end = index-1
            index = partition(data_list,begin,index-1)
        else:
            begin = index+1
            index = partition(data_list,index+1,end)
    return data_list[index]```

執行一下看看效果：

​```python
    data_list = [25, 77, 52, 49, 85, 28, 1, 28, 100, 36]
    print(data_list)
    for i in range(1,11):
        print(f"第 {i} 大元素是 {find_top_k(data_list,i)}")```
執行結果如下所示

​```python
[25, 77, 52, 49, 85, 28, 1, 28, 100, 36]
第 1 大元素是 100
第 2 大元素是 85
第 3 大元素是 77
第 4 大元素是 52
第 5 大元素是 49
第 6 大元素是 36
第 7 大元素是 28
第 8 大元素是 28
第 9 大元素是 25
第 10 大元素是 1

下面解釋一下為什么時間復雜度是 O(n):

第一次分區查找，我們需要對大小為 n 的數組執行分區操作，需要遍歷 n 個元素。第二次分區查找，我們只需要對大小為 n/2 的數組執行分區操作，需要遍歷 n/2 個元素。依次類推，分區遍歷元素的個數分別為、n/2、n/4、n/8、n/16.…… 直到區間縮小為 1。 如果我們把每次分區遍歷的元素個數加起來，就是：n+n/2+n/4+n/8+…+1。這是一個等比數列求和，最后的和等於 2n-1。所以，上述解決思路的時間復雜度就為 O(n)。

小結

快速排序和歸並排序都是分治的思想，代碼都通過遞歸來實現，歸並排序的重點在於 merge 函數，而快排的重點在於 partition 函數。歸並排序優點: 任何情況下時間復雜度穩定在 O(nlogn), 缺點：不是原地排序算法，需要額外的內存空間。快速排序是一種原地排序算法，平均時間復雜度為 O(nlogn)，但極端情況時間復雜度會退化成 O(n^2)，雖然這種情況的概率非常小，仍需要合理的選擇分區鍵，避免左右分區極度不平衡。