题解 | 查找第K小数

#include <iostream>
#include <vector>
#include <unordered_set>
#include <algorithm>
using namespace std;

// 分区操作
int partition(vector<int>& vec, int left, int right) {
    int pivot = vec[left];
    int i = left + 1, j = right;

    while (i <= j) {
        while (i <= j && vec[i] <= pivot) i++; // 找到第一个大于 pivot 的元素
        while (i <= j &&
                vec[j] > pivot) j--; // 找到第一个小于等于 pivot 的元素
        if (i < j) {
            swap(vec[i], vec[j]); // 交换这两个元素
            i++, j--;
        }
    }
    swap(vec[left], vec[j]); // 将基准值放到正确的位置
    return j; // 返回基准值的索引
}

// 快速选择算法
int findKthSmallest(vector<int>& vec, int left, int right, int k) {
    if (left == right) {
        return vec[left]; // 如果区间只有一个元素，直接返回
    }

    // 分区操作
    int pivotIndex = partition(vec, left, right);

    // 判断是否找到第 k 小的元素
    if (pivotIndex == k - 1) {
        return vec[pivotIndex];
    } else if (pivotIndex < k - 1) {
        return findKthSmallest(vec, pivotIndex + 1, right,
                               k); // 在右半部分递归查找
    } else {
        return findKthSmallest(vec, left, pivotIndex - 1,
                               k); // 在左半部分递归查找
    }
}

int main() {
    int n;
    cin >> n;
    unordered_set<int> uset;
    int num;
    for (int i = 0; i < n; i++) {
        cin >> num;
        uset.insert(num);
    }
    cin >> num;

    // 去重并排序
    vector<int> vec(uset.begin(), uset.end());

    // 输出第 k 小的数
    cout << findKthSmallest(vec, 0, vec.size() - 1, num) << endl;

    return 0;
}