[Algorithm] Quick Selection

 

 

 

소스코드

#include <cstdio>
#include <algorithm>

using namespace std;

const int MX = 5000001;

int partition(int* a, int left, int right)
{
    int pivot_idx = right;
    int pivot_val = a[pivot_idx];
    
    swap(a[pivot_idx], a[right]);
    
    int cnt = left;
    for (int i = left; i < right; ++i) {
        if (a[i] <= pivot_val) { //important
            swap(a[i], a[cnt++]);
        }
    }
    swap(a[right], a[cnt]);
    return cnt;
}

// return kth smallest number
int kth_elem(int* a, int l, int r, int k)
{
    int pivot = partition(a, l, r);
    if (pivot == k) {
        return a[pivot];
    }
    else if (pivot < k) {
        return kth_elem(a, pivot + 1, r, k);
    }
    else {
        return kth_elem(a, l, pivot - 1, k);
    }
}

int main()
{
    int n, k, a[MX];
    scanf("%d %d", &n, &k);
    for (int i = 0; i < n; ++i)
        scanf("%d", &a[i]);

    printf("%d\n", kth_elem(a, 0, n - 1, k - 1));
    return 0;
}

 

 

 

 

 

공부한 것을 정리하기 위한 기록입니다.

틀린 부분이 있다면 부드럽게 알려주세요.