Kth smallest element, C++ (gcc)

Kth smallest element

const int SIZE = 9;

inline int get_val(const vector<int>& v){
    return v[0]*v[0] + v[1]*v[1];
}

bool cmp(const vector<int>& v1, const vector<int>& v2){
    return get_val(v1) < get_val(v2);
}

class Solution {
    int median(int l, int r, vector<vector<int>>& points){
        if((r - l + 1) < SIZE){
            sort(points.begin()+l, points.begin()+r+1, cmp);
            return get_val(points[(l+r)/2]);
        }
        
        int i = l;
        for(int j = l; j <= r; j += SIZE){
            int _l = j;
            int _r = (j+SIZE-1 <= r)? j+SIZE-1 : r;
            sort(points.begin()+_l, points.begin()+_r+1, cmp);
            swap(points[i], points[(_l+_r)/2]);
            i ++;
        }
        
        return median(l, i-1, points);
    }
    
    int get_kth(int l, int r, vector<vector<int>>& points, int K){
        if(l == r) return get_val(points[l]);
        
        int md = median(l, r, points);
        int i = l, j = r, _r = r;
        while(_r >= l && get_val(points[_r]) == md) _r --;
        while(i <= _r){
            if(get_val(points[i]) == md){
                swap(points[_r--], points[i]);
                while(_r >= l && get_val(points[_r]) == md) _r --;
            }
            i ++;
        }
        // [l, i] doesn't contains md
        i = l, j = _r;
        while(i <= j){
            while(i <= _r && get_val(points[i]) < md) i ++;
            while(j >=  l && get_val(points[j]) > md) j --;
            if(i > j) break;
            swap(points[i], points[j]);
        }
        
        // cout << i << " " << j << " " << _r << " " << r << endl;
        int lft = i - l;
        if(lft >= K) return get_kth(l, i-1, points, K);
        K -= lft;
        if((r - _r) >= K) return md;
        K -= (r - _r);
        return get_kth(i, _r, points, K);
    }
    
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        int n = points.size();
        int m = get_kth(0, n-1, points, K);
        cout << m << endl;
        vector<vector<int>> ans;
        for(int j = 0; j < n && ans.size() < K; j ++)
            if(get_val(points[j]) < m) ans.push_back(points[j]);
        for(int j = 0; j < n && ans.size() < K; j ++)
            if(get_val(points[j]) == m) ans.push_back(points[j]);
        return ans;
    }
};


run \| edit \| history \| help	0