归并排序和快速排序的衍生问题（二）

取数组中第n大的元素

取数组中第n大的元素这个问题，相信大家在学习数据结构这门课中都遇到过。通常我们会使用某一排序算法先将数组排序，然后在来找数组中第n大的元素。此时，数组中第n大的元素其在数组中的位置也为第n个。

利用这一思想，我们可以使用快速排序算法来将该求解算法的时间复杂度变为O(n)。为什么呢？在快速排序算法中，每一趟快速排序都会确定一个元素（基准）的位置，例如：假设现在基准为5，经过一趟快速排序后，我们知道在第五个元素（即该元素为5）之前的元素都小于5，在第五个元素之后的元素都大于5，这时我们不就确定了元素（基准）为5在数组中的位置吗？

我们先使用第一种方法来求解数组中第n大的元素，即先排序再找第n大的元素，其基本代码如下：

// 3-路快速排序算法
void __quickSort(int arr[], int l, int r) {

    if (l >= r)
        return;

    swap(arr[l], arr[rand() % (r - l + 1) + l]);
    int v = arr[l];

    // arr[l+1...lt] < v
    int lt = l;
    // arr[gt...r] > v
    int gt = r + 1;
    // arr[lt+1...i) == v
    int i = l + 1;

    while (i < gt) {
        if (arr[i] < v) {
            swap(arr[i], arr[lt + 1]);
            lt++;
            i++;
        } else if (arr[i] > v) {
            swap(arr[i], arr[gt - 1]);
            gt--;
        } else {
            i++;
        }
    }
    swap(arr[l], arr[lt]);

    __quickSort(arr, l, lt - 1);
    __quickSort(arr, gt, r);
}

int __selection(int arr[], int n, int k) {

    srand(time(NULL));
    __quickSort(arr, 0, n - 1);

    return arr[k];
}

// 寻找arr数组中第k大的元素
int selection(int arr[], int n, int k) {

    // 如果k >= 0 && k < n不成立，则终止程序执行
    assert(k >= 0 && k < n);

    return __selection(arr, n, k);
}

现在，我们来修改一下代码，使其成为第二种方法，其代码如下：

int __partition(int arr[], int l, int r) {

    swap(arr[l], arr[rand() % (r - l + 1) + l]);
    int v = arr[l];

    // arr[l+1...lt] < v
    int lt = l;
    // arr[gt...r] > v
    int gt = r + 1;
    // arr[lt+1...i) == v
    int i = l + 1;

    while (i < gt) {
        if (arr[i] < v) {
            swap(arr[i], arr[lt + 1]);
            lt++;
            i++;
        } else if (arr[i] > v) {
            swap(arr[i], arr[gt - 1]);
            gt--;
        } else {
            i++;
        }
    }
    swap(arr[l], arr[lt]);

    return lt;
}

int __selection(int arr[], int l, int r, int k) {

    if (l == r) return arr[l];

    int p = __partition(arr, l, r);

    if (k == p) return arr[p];
    else if (k < p)
        return __selection(arr, l, p - 1, k);
    else
        return __selection(arr, p + 1, r, k);
}

// 寻找arr数组中第k大的元素
int selection(int arr[], int n, int k) {

    // 如果k >= 0 && k < n不成立，则终止程序执行
    assert(k >= 0 && k < n);
    srand(time(NULL));

    return __selection(arr, 0, n - 1, k);
}

这样，我们就解决了求取数组中第n个大元素的问题。

注：在main()中的代码如下：

int main() {

    // 生成一个大小为n, 包含0...n-1这n个元素的随机数组arr
    int n = 10000;
    int *arr = TestHelper::generateOrderedArray(n);
    TestHelper::shuffleArray(arr, n);

    // 验证selection算法, 对arr数组求第i小元素, 应该为i
    for (int i = 0; i < n; i++) {
        assert(selection(arr, n, i) == i);
        cout << "test " << i << " complete." << endl;
    }
    cout << "Test selection completed." << endl;

    delete[] arr;
    return 0;
}

由于测试方法为遍历输出，故不贴出程序运行结果。