neko-suki / icpc-library

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

quick select

neko-suki opened this issue · comments

    // [l, r)
    // divided into [left, pos), [pos, r)
    // arr[pivot] <= arr[pos] 
    int partition(vector<int>& arr, int l, int r, int pivot){
        swap(arr[r-1], arr[pivot]);
        int pos = l;
        for(int i = l;i < r-1;i++){
            if (arr[i] < arr[r-1]){
                swap(arr[i], arr[pos]);
                pos++;
            }
        }
        swap(arr[pos], arr[r-1]);
        return pos;
    }
    
    // [l,r)
    // k is 0 based
    int search(vector<int> & arr, int l, int r, int k){
        if (l+1 == r){
            return arr[l];     
        }
        int pivot = (l+r)/2; // random should be better choice
        pivot = partition(arr, l, r, pivot);
        if (pivot == k){
            return arr[pivot];
        }
        if (pivot < k){
            return search(arr, pivot+1, r, k);   
        } else {
            return search(arr, l, pivot, k);
        }
    }