Kth smallest, C++ (gcc)

Kth smallest

#include <bits/stdc++.h>

using namespace std;

int median(int l, int r, vector<int>& v){
    if((r - l + 1) <= 5){
        sort(v.begin(), v.end());
        int mid = (r - l)/2 + l;
        return v[mid];
    }
    int len = r - l + 1;
    int each = len / 5;
    
    vector<int> m;
    for(int j = l; j <= r; j += 5){
        int st = j, ed = (j+5-1 <= r)? (j + 5-1) : r;
        sort(v.begin()+st, v.begin()+ed+1);
        
        int x = (j+5-1 <= r)? j+2 : (j + (ed - st)/2);
        m.push_back(v[x]);
    }
    
    return median(0, (int)m.size() - 1, m);
}

int foo(int l, int r, vector<int>& v, int k){
    if(l == r) return v[l];
    int m = median(l, r, v);
    
    int i = l, j = r;
    while(i <= j){
        while(v[i] <= m) i ++;
        while(v[j] > m) j --;
        if(i > j) break;
        
        swap(v[i], v[j]);
        i ++, j --;
    }
    
    // for(int j = l; j <= r; j++) cout << v[j] << " "; cout << " --> " << m << endl;
    
    int cnt_less_m = j - l + 1;
    if(cnt_less_m >= k) return foo(l, j, v, k);
    return foo(i, r, v, k-cnt_less_m);
}

bool check(int x, const vector<int>& A, const int B){
    int cnt = 0;
    for(auto i : A) 
        if(i <= x) cnt ++;
    return (cnt >= B);
}

int foo(int l, int r, const int B, const vector<int>& A){
    if(l == r) return l;
    int m = (l+r) >> 1;
    if(check(m, A, B)) return foo(l, m, B, A);
    return foo(m+1, r, B, A);
}

int Solution(const vector<int> &A, int B) {
    return foo(0, 1e5+7, B, A);
}

int main(){
    ios_base :: sync_with_stdio(0);
    cin.tie(0);

int t;
    cin >> t;
    while(t--){
        int n;
        cin >> n;
        vector<int> v(n);
        for(int j = 0; j < n; j++) cin >> v[j];
        int k;
        cin >> k;
        // cout << foo(0, n-1, v, k) << '\n';
        cout << Solution(v, k) << '\n';
    }
    return 0;
}


run \| edit \| history \| help	0