如何用快排思想在O(n)內查找第K大元素--王爭《數據結構和算法之美》

前言

半年前在極客時間訂閱了王爭的《數據結構和算法之美》，現在決定認真去看看。看到如何用快排思想在O(n)內查找第K大元素這一章節時發現王爭對歸並和快排的理解非常透徹，講得也非常好，所以想記錄總結一下。文章內容主要分析歸並排序和快速排序原理，並根據它們共同的分治思想，引出如何在 O(n) 的時間復雜度內查找一個無序數組中的第 K 大元素？

歸並排序原理

核心思想：將數組從中間分成前后兩部分，然后對前后兩部分分別進行排序，再將排序好的兩個部分有序合並在一起，這樣整個數組有序。

歸並排序使用的就是分治思想。分治，顧名思義，就是分而治之，講一個大的問題分解成小的問題來解決，小的問題解決了大的問題也就解決了。分治算法一般都是用遞歸來實現，分治是一種解決問題的處理思想，遞歸是一種編程技巧，兩者並不沖突。以下重點討論如何用遞歸代碼來實現歸並排序。下面是歸並排序的遞推公式。

遞推公式：
merge_sort(p...r) = merge(merge_sort(p...q), merge_sort(q+1...r))

終止條件：
p >= r 不用繼續分解

具體解釋如下：

merge_sort(p...r) 表示給下標從 p 到 r 之間的數組排序。將這個排序問題轉化為兩個子問題 merge_sort(p...q) 和merge_sort(q+1...r)，其中 q 為 p 和 r 的中間位置，即(p+r)/2。當前后兩個子數組排好序之后，再將它們合並在一起，這樣下標從 p 到 r 之間的數據也就排序好了。

C語言代碼實現：

// 歸並排序算法, A 是數組，n 表示數組大小
void mergeSort(int *a, int n){
    mergeSortC(a, 0, n-1);
}

// 遞歸調用函數
void mergeSortC(int *a, int left, int right){
    // 遞歸終止條件
    if (left >= right)
        return;

    int mid = left + (right - left)/2;
    mergeSortC(a, left, mid);
    mergeSortC(a, mid+1, right);
    merge(a, left, mid, right);
}

// 合並函數
void merge(int *a, int left, int mid, int right){
    int i = left, j = mid+1, k = 0;
    int *tmp = new int[right-left+1];  // 申請一個大小為right-left+1臨時數組
    while (i <= mid && j <= right){
        if(a[i] < a[j])
            tmp[k++] = a[i++];
        else
            tmp[k++] = a[j++];
    }

    while (i <= mid)
        tmp[k++] = a[i++];

    while (j <= right)
        tmp[k++] = a[j++];

    for (i=0; i <= right-left; i++){
        a[left+i] = tmp[i];
    }

    delete[] tmp;
}

歸並排序的時間復雜度任何情況下都是 O(nlogn)，看起來非常優秀（快速排序最壞情況系時間復雜度也是 O(n2)）。但歸並排序並沒有像快排那樣應用廣泛，因為它有一個致命的“弱點”，那就是歸並排序不是原地排序算法。原因是合並函數需要借助額外的存儲空間，空間復雜度為 O(n)。

C++實現：

void merge(std::vector<int>& a, int left, int mid, int right) {
    int i = left;
    int j = mid + 1;
    int k = 0;
    std::vector<int> v(right - left + 1);
    while (i <= mid && j <= right) {
       v[k++] = a[i] < a[j] ? a[i++] : a[j++];
    }
    while (i <= mid) {
        v[k++] = a[i++];
    }
    while (j <= right) {
        v[k++] = a[j++];
    }
    for (i = 0; i < v.size(); ++i) {
        a[left + i] = v[i];
    }
}

void mergeSort(std::vector<int>& a, int left, int right) {
    if (left >= right) return;
    int mid = left + (right - left) / 2;
    mergeSort(a, left, mid);
    mergeSort(a, mid+1, right);
    merge(a, left, mid, right);
}

void mergeSort(std::vector<int>& a) {
    mergeSort(a, 0, a.size() - 1);
}

快速排序原理

核心思想：選取一個基准元素（pivot，比 pivot 小的放到左邊，比 pivot 大的放到右邊，對 pivot 左右兩邊的序列遞歸進行以上操作。

快速排序也是根據分治、遞歸的處理思想實現。地推公式如下：

遞推公式：
quick_sort(p…r) = quick_sort(p…q-1) + quick_sort(q+1...r)

終止條件：
p >= r

C語言代碼實現：

// 快速排序算法, A 是數組，n 表示數組大小
void quickSort(int *a, int n){
    quickSortC(a, 0, n-1);
}

// 快排遞歸函數
void quickSortC(int *a, int left, int right){
    // 遞歸終止條件
    if (left >= right)
        return;
    // 獲取分區點
    int pivot = partition(a, left, right);
    quickSortC(a, left, pivot-1);
    quickSortC(a, pivot+1, right);
}

/* 原地分區函數，非常巧妙，以a[right]為基准，運算結果
 * 是i前面的元素都小於pivot，i后面的元素大於等於pivot */
int partition(int *a, int left, int right){
    int pivot = a[right];
    int i = left;
    for (int j=left; j < right; j++){
        if (a[j] < pivot){
            swap(a[i], a[j]);
            i++;
        }
    }
    swap(a[i], a[right]);
    return i;
}

 快速排序的算法的平均時間復雜度是 O(nlogn)，最壞時間復雜度是 O(n2)，空間復雜度是O(1)。快速排序不是一個穩定的排序算法。

C++ 實現：

int partition(std::vector<int>& a, int left, int right) {
    using std::swap;
    int pivot = a[right];
    int j = left;
    for (int i = left; i < right; ++i) {
        if (a[i] < pivot)
            swap(a[i], a[j++]);
    }
    swap(a[right], a[j]);
    return j;
}

void quickSort(std::vector<int>& a, int left, int right) {
    if (left >= right) return;
    int pivot = partition(a, left, right);
    quickSort(a, left, pivot-1);
    quickSort(a, pivot+1, right);
}

void quickSort(std::vector<int>& a) {
    quickSort(a, 0, a.size() - 1);
}

歸並排序和快速排序的區別

快排和歸並用的都是分治思想，遞歸公式和代碼都非常相似，但它們的區別在哪里呢？

 

 由上圖可以發現，歸並排序的處理過程是由下到上的，先處理子問題，然后合並。而快排正好相反，其處理過程是由上而下的，先分區，然后處理子問題。歸並排序雖然是穩定的，時間復雜度是 O(nlogn)的排序算法，但它是非原地排序算法。快排通過設計巧妙的原地分區函數，可以實現原地排序，解決歸並排序占用太多內存的問題。

第 K 大元素

快排核心思想就是分治和分區，我們可以利用分區的思想來求解開篇問題： O(n)時間復雜度內求無序數組中的第 K 大元素。

 C語言代碼實現：

// top K 算法， A 是數組，n 表示數組大小，k 表示第 k 大
int getTopK(int *a, int n, int k){
    if (a == nullptr || n < k)
        return -1;
    
    return topK(a, 0, n-1, k);
}

int topK(int *a, int left, int right, int k){
    int p = partition(a, left, right);
    if (k == p+1)
        return a[p];

    if(k < p+1)
        return topK(a, left, p-1, k);
    else
        return topK(a, p+1, right, k);
}

/* 原地分區函數，非常巧妙，以a[right]為基准，運算結果
 * 是i前面的元素都大於pivot，i后面的元素小於於等於pivot */
int partition(int *a, int left, int right){
    int pivot = a[right];
    int i = left;
    for (int j=left; j < right; j++){
        if (a[j] > pivot){
            swap(a[i], a[j]);
            i++;
        }
    }
    swap(a[i], a[right]);
    return i;
}

LeetCode 215 C++實現：

class Solution {
public:
    int partition(vector<int>& nums, int left, int right) {
        using std::swap;
        int pivot = nums[right];
        int j = left;
        for (int i = left; i < right; ++i) {
            if (nums[i] > pivot) 
                swap(nums[i], nums[j++]);
        }
        swap(nums[right], nums[j]);
        return j;
    }
    int getTopK(vector<int>& nums, int left, int right, int k) {
        if (left >= right) return nums[left];
        int pivot = partition(nums, left, right);
        if (pivot + 1 == k) 
            return nums[pivot];
        return (pivot + 1 < k) ? getTopK(nums, pivot+1, right, k) 
                               : getTopK(nums, left, pivot-1, k);
    }
    int findKthLargest(vector<int>& nums, int k) {
        return getTopK(nums, 0, nums.size() - 1, k);
    }
};