백준 11004 K번째 수

문제)

수 N개 A1, A2, ..., AN이 주어진다. A를 오름차순 정렬했을 때, 앞에서부터 K번째 있는 수를 구하는 프로그램을 작성하시오.

입력)
첫째 줄에 N(1 ≤ N ≤ 5,000,000)과 K (1 ≤ K ≤ N)이 주어진다.

둘째에는 A1, A2, ..., AN이 주어진다. (-109 ≤ Ai ≤ 109)

출력)
A를 정렬했을 때, 앞에서부터 K번째 있는 수를 출력한다.

sol)

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;


void quickSort(vector<int> &A, int S, int E, int K);
int partition(vector<int> &A, int S, int E);
void swap(vector<int>& A, int S, int E);

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    int N, K;
    cin >> N >> K;

    vector<int> A(N, 0);

    for (int i = 0; i < N;i++) {
        cin >> A[i];

    }
    //0 = 시작 / N-1 종료
    quickSort(A, 0, N-1, K-1);
    cout << A[K - 1];

    return 0;
}

void quickSort(vector<int>& A, int S, int E, int K) {
    int pivot = partition(A, S, E);
    if (pivot == K) {
        return;

    }
    else if (K < pivot) {
        quickSort(A, S, pivot - 1, K);
    }
    else {
        quickSort(A, pivot + 1, E , K);
    }

}

int partition(vector<int>& A, int S, int E) {
    if (S + 1 == E) {
        if (A[S] > A[E]) swap(A, S, E);
        return E;
    }
    int M = (S + E) / 2;
    swap(A, S, M);
    int pivot = A[S];
    int i = S + 1;
    int j = E;

    while (i <= j) {
        while (pivot < A[j] && j > 0) j--;
        while (pivot < A[i] && i < A.size() - 1) i++;
        if (i <= j) swap(A, i++, j--);
    }
    A[S] = A[j];
    A[j] = pivot;
    return j;
}
void swap(vector<int>& A, int S, int E) {
    int temp = A[S];
    A[S] = A[E];
    A[E] = temp;
}