[백준 16455] K번째 수 찾는 함수

https://www.acmicpc.net/problem/16455

 

 

처음 이 문제를 봤을 때 단순히 피벗 정렬을 이용했는데 데이터 안에 최악의 피벗 정렬이 되게끔 만드는 데이터가 있었다. 피벗정렬의 최악 시간복잡도는 O(N²)으로 메모리초과가 나거나 시간 초과가 무조건 나는 상황이었다..

 

해결법을 찾아봤는데 데이터를 랜덤으로 shuffle 후 nth_element을 사용하는 방법이 있었는데 너무 간결하게 끝나서 또다른 방법을 찾아보기 시작했다.

 

힌트를 얻어 binary_search + radix sort를 이용하는 방법을 공유받았는데 기본적으로 radix sort는 보통 10으로 했지만 이런 경우는 수의 크기가 작은 경우라 아이디어를 떠올리기 쉽지 않았다.

 

해당 문제는 수의 범위가 -10억 ~ 10억으로 크기가 워낙 크니 해당 수를 제곱근으로 나눠 데이터를 관리하였다. 20억의 제곱근은 대략 44721이고, 버킷의 크기를 오차가 없이 관리하려고 45000으로 고정하였다.

 

그리고 배열에 넣기 위해 입력되는 값에 10억을 더해 음수를 없앴으며 버킷으로 나눈 몫/나머지를 카운트 할 2차원 vector 두개가 필요하다. 몫을 가지는 벡터를 cA[45000], 나머지를 가지는 벡터를 cB[45000]로 두었다.

 

예를 들어 857436이라는 입력 데이터가 있으면 10억을 더해 1000857436이라는 수가 되고 해당 수를 45000으로 나누면 몫이 22241, 나머지가 12436이 된다. 누적 합으로 구하기 때문에 해당 수를 직접 담는게 아닌 해당 버킷에 수가 있다는 걸 카운트만 해준다.

 

처음에는 k번째 수의 대략적인 버킷 위치를 찾기 위한 몫 기준으로 기수 정렬을 수행해주고 각 벡터를 누적 합으로 만들어준다. 누적 합 벡터는 항상 오름차순으로 정렬되어 있기 때문에 k번째 수가 누적 합 vector의 몇 번째에 있는지 이분 탐색을 이용해 찾으면 된다. 몫에 대한 이분 탐색을 이용해 몇번째 버킷에 있는지 찾으면 해당 버킷에 해당하지 않는 수들은 전부 정렬 대상에 포함되지 않는다. 수가 들어있는 리스트를 다시 한번 순회해 45000 * 버킷 위치를 벗어나는 데이터는 전부 스킵하고 해당 버킷에 포함되는 수들 중 45000의 나머지를 이용해 다시 한번 기수 정렬을 하고 누적 합 벡터를 만든다.

 

k번째 수를 몫에 대해 찾았으니 해당 버킷 보다 작은 수의 갯수를 빼 주고 나머지에 대한 부분도 이분 탐색을 해서 해당 버킷보다 작은 수의 갯수를 찾은 다음 버킷 * 몫 + 나머지 - 10억을 해 원래대로 돌려주면 정답이 된다.

 

#include <bits/stdc++.h>

using namespace std;

int bucket = 45000;
int b_search(vector<int> &a, int target) {
    int s = 0;
    int e = bucket;
    int res = 0;
    while(s <= e) {
        int mid = (s + e) / 2;
        if(a[mid] >= target) {
            res = mid;
            e = mid - 1;
        } else {
            s = mid + 1;
        }
    }
    return res;
}
int kth(vector<int> &a, int k) {
    vector<int> cA(bucket);
    vector<int> cB(bucket);

    for(int & i : a) {
        i += 1000000000;
        cA[i / bucket]++;
    }
    // bucket + radix sort
    for(int i = 1; i < cA.size(); i++) {
        cA[i] += cA[i - 1];
    }
    int b = b_search(cA, k);
    for(int & i : a) {
        if(i / bucket != b) {    // 범위 밖 값들은 정렬할 필요가 없음
            continue;
        }
        cB[i % bucket]++;
    }
    for(int i = 1; i < cB.size(); i++) {
        cB[i] += cB[i - 1];
    }
    int t = k - (b == 0 ? 0 : cA[b - 1]);
    int d = b_search(cB, t);
    return b * bucket + d - 1000000000;
}