O(n)寻找数列第K小

期末考完了，博客继续更新。

题目链接点这儿

O(n)求数列第K大和O(n)求数列第K小的思路一样，利用分治求解。

算法思路很简单，我们从数列里挑一个数，然后看比它大和比它小的元素各有多少。如果比它小的恰为k-1个，那么这个元素就是第K小，若小于，那么就从比它大的数里找第k-x-1小（x为比它小的元素个数）。若大于，则在比它小的数里找第k小。

这个算法的均摊复杂度为O(n)。

当然我们也可以用algorithm里的nth_element()函数，详细看这里。

过题代码

#include <bits/stdc++.h>

#define up(i, lower, upper) for(int i = lower; i < upper; i++)
#define down(i, lower, upper) for(int i = upper-1; i >= lower; i--)

using namespace std;
typedef pair<int, int> pii;
typedef pair<double, double> pdd;
typedef vector<int> vi;
typedef vector<pii> vpii;
typedef long long ll;
typedef unsigned long long ull;

const double pi = acos(-1.0);
const double eps = 1.0e-9;

template<class T>

inline bool read(T &n){
    T x = 0, tmp = 1; char c = getchar();
    while((c < '0' || c > '9') && c != '-' && c != EOF) c = getchar();
    if(c == EOF) return false;
    if(c == '-') c = getchar(), tmp = -1;
    while(c >= '0' && c <= '9') x *= 10, x += (c - '0'),c = getchar();
    n = x*tmp;
    return true;
}

template <class T>
inline void write(T n) {
    if(n < 0) {
        putchar('-');
        n = -n;
    }
    int len = 0,data[20];
    while(n) {
        data[len++] = n%10;
        n /= 10;
    }
    if(!len) data[len++] = 0;
    while(len--) putchar(data[len]+48);
}

///---------------------------------------------------------
int a[5000010];

int find_kth(int l, int r, int k) {
    int random = l + rand()%(r - l + 1);
    swap(a[random], a[l]);
    int m = l, cnt = 0;
    up(i, l+1, r+1) {
        if(a[i] < a[l])
            swap(a[++m], a[i]), cnt++;
    }
    swap(a[m], a[l]);
    if(cnt == k-1) return m;
    else if(cnt < k-1) return find_kth(m+1, r, k - cnt - 1);
    else return find_kth(l, m-1, k);
}

int main() {
    int n, k;
    read(n), read(k);
    k = n - k + 1;
    up(i, 0, n) read(a[i]);
    write(a[find_kth(0, n-1, k)]);
    puts("");
    return 0;
}