백준 16455번: K번째 수 찾는 함수 (Java)

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

16455번: K번째 수 찾는 함수

 

 

 

 

 

 

 

 

---

 

풀이

N <= 500만이라 O(NlogN)은 안 됩니다. 실제로 O(NlogN) 정렬(종류 상관 x)을 수행한 후 K번째 원소의 값을 출력하는 건 시간초과가 뜨기에, 다른 방법을 찾아봐야 합니다.

 

K번째 수를 찾는 방법에는 quick select라는 것이 있다고 합니다. 대충 설명하자면, quicksort 방식과 마찬가지로 pivot을 잡아 원소를 3개의 집합 

L := {x | x < pivot}

C := {x | x == pivot}

R := {x | x >= pivot}

으로 나누어, |L| + |C| < K이면 R, |L| > K이면 L에서 새로운 pivot을 찾고,

|L| <= K or |L| + |C| >= K이면 적당히 L or C에서 K번째 원소를 찾는 방법입니다.

quicksort랑 다르게 amortized O(N) 시간복잡도를 가지지만, worst case는 똑같이 O(N^2)입니다. 따라서, worst-case를 O(N)에 처리할 수 있는 median of medians 알고리즘이 필요하나, 실용성도 떨어지고, 사실 좀 구현하기 귀찮습니다. 그래서, 백준 ID 'bubbler'님 풀이를 참고하여 풀었습니다.

 

(주어진 수) + 10억을 pq + r (0 <= r < p) 꼴로 나타냅니다. 그러면, (주어진 수) <-> (q, r) 간의 bijection이 생깁니다. 이를 이용하여 풀이를 해봅시다. '몫을 저장하는 배열' arr1, '나머지를 저장하는 배열' arr2 총 2개를 만들어,

 

1) N개의 수를 p로 나눈 몫 q를 구하여, 'q에 대응되는 숫자가 1개 있다'는 의미로 arr1[q]의 값을 1 늘려줍니다.

2) 이후, 1번째 index부터 arr1의 값을 더하여, 합이 K 이상이 될 경우(== 해당 index 이하에 대응되는 숫자가 K개 이상이다) 해당 인덱스 x를 기억해둡니다(주: 코드에서는 p로 써서 보기 헷갈리지도..?).

3) N개의 수를 다시 순회하여, p로 나눈 몫이 x인 수의 나머지 r를 구하여 arr2[r]++을 해줍니다.

4) 2)와 동일하게, 세면서 k번째 이상이 되면 멈춥니다. 그 때의 인덱스 r을 기억해둡니다.

정답은 x * p + r - 10000000이 됩니다.

 

숫자의 범위를 R이라고 했을 때(본 문제: R=20억)이 알고리즘은 O(max(N, p, R/p))가 됩니다. p = ceil(sqrt(20억))으로 잡으면 최소화할 수 있으며, 저는 코드의 가독성을 살짝 생각하여 p = 50000으로 잡았습니다. 

 

이 알고리즘을 일컫는 이름이 있을텐데, 백준 코드를 보고 배워서 뭔지는 모르겠네요.

 

그리고 C++을 쓴다면 nth_element?인가하는 걸 쓰면 1줄에 풀립니다.

 

 

코드

public class Test {
    static final int INF = 1_000_000_000;
    
    int kth(int[] a, int k) {
        // 정답: (p - 20000) * 50000 + r
        int[] arr1 = new int[50000];
        int[] arr2 = new int[50000];
        for (int num: a) {
            arr1[(num+INF)/50000]++; 
        }
        
        int p = 0;
        for (int i = 0; i < 50000; i++) {
            k -= arr1[i];
            if (k <= 0) {
                k += arr1[i];
                p = i;
                break;
            }
        }
        
        for (int num: a) {
            if ((num+INF)/50000 == p) {
                arr2[(num+INF)%50000]++;
            }
        }
        
        int r = 0;
        for (int i = 0; i < 50000; i++) {
            k -= arr2[i];
            if (k <= 0) {
                r = i;
                break;
            }
        }
        return (p - 20000) * 50000 + r;
    }
}