跳转至

kth_element()

之前学STL道德时候见过这个函数,不过没有了解其原理,今天把他的实现代码自己过一遍,应该能掌握更好一点

// 用来选取首、中、尾的中位数
int median3(vector<int>& nums, int l, int r) {
    int mid = l + (r - l) / 2;
    if (nums[l] > nums[mid]) swap(nums[l], nums[mid]);
    if (nums[l] > nums[r]) swap(nums[l], nums[r]);
    if (nums[mid] > nums[r]) swap(nums[mid], nums[r]);
    return mid;
}

int quickselect(vector<int>& nums, int l, int r, int k) {
    // 选基准(首、中、尾的中位数)
    int mid = midian3(nums, l, r);
    swap(nums[mid], nums[r]);
    int pivot = nums[r];
    // 分区
    int pos = l;
    for (int i = l; i < r; i ++) {
        if (nums[i] < pivot) swap(nums[i], nums[pos ++]);
    }
    swap(nums[pos], nums[r]);

    // 单边递归
    if (pos == k) return nums[pos];
    else if (pos < k) return quickselect(nums, pos + 1, r, k);
    else return quickselect(nums, l, pos - 1, k);
}

时间复杂度是\(O(n + \frac{n}{2} + \frac{n}{4} + ... + 1) = O(2n)\)