문제
N개의 수가 주어졌을 때 오름차순으로 정렬한 결과에서 K번째 수를 구하는 문제다. N은 최대 5,000,000으로 빠른 I/O와 효율적인 정렬이 필요하다.
입력
- 첫째 줄: N, K
- 둘째 줄: N개의 수
출력
오름차순 정렬 시 K번째 수
예제
| 입력 | 출력 |
|---|---|
5 2 4 1 2 3 5 | 2 |
풀이
퀵 셀렉트(Quick Select) 알고리즘으로 전체를 정렬하지 않고 K번째 원소만 찾는다. 파티션 후 피벗의 위치가 K-1이면 바로 종료하고, 그렇지 않으면 K가 속하는 절반만 재귀 탐색한다.
- N, K를 입력받고 배열에 N개의 수를 저장한다.
quickSort(arr, 0, N-1, K)를 호출한다.partition에서 중간값을 피벗으로 선택해 left에 두고 양방향 포인터로 분할한다.- 피벗 인덱스
pi + 1 == K이면 탐색 종료,pi + 1 < K이면 오른쪽만, 그렇지 않으면 왼쪽만 재귀한다. arr[K-1]을 출력한다.
핵심 아이디어: K번째 원소를 찾는 데 전체 정렬(O(N log N))은 불필요하다. 퀵 셀렉트는 평균 O(N)에 K번째 원소를 찾는다. 피벗 위치가 K-1과 일치하는 순간 나머지 원소 정렬 없이 즉시 답을 반환할 수 있다.
코드
package com.ssafy.an.day049;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;
public class Day07BOJ11004Fail1K번째수 { // 11004 pivot sort
public static void main(String[] args) throws Exception {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());
int N = Integer.parseInt(st.nextToken());
int K = Integer.parseInt(st.nextToken());
int[] arr = new int[N];
st = new StringTokenizer(br.readLine());
for (int n = 0; n < N; n++) {
arr[n] = Integer.parseInt(st.nextToken());
}
quickSort(arr, 0, N - 1, K);
System.out.println(arr[K - 1]);
}
private static void quickSort(int[] arr, int left, int right, int length) {
if (left >= right)
return;
int pi = partition(arr, left, right);
if (pi + 1 == length)
return;
else if (pi + 1 < length)
quickSort(arr, pi + 1, right, length);
else
quickSort(arr, left, pi - 1, length);
}
private static int partition(int[] arr, int left, int right) {
int mid = (left + right) / 2;
swap(arr, left, mid);
int pivot = arr[left];
int i = left, j = right;
while (i < j) {
while (pivot < arr[j]) {
j--;
}
while (i < j && pivot >= arr[i]) {
i++;
}
swap(arr, i, j);
}
arr[left] = arr[i];
arr[i] = pivot;
return i;
}
private static void swap(int[] arr, int i, int j) {
int temp = arr[i];
arr[i] = arr[j];
arr[j] = temp;
}
}복잡도
- 시간: 평균 O(N), 최악 O(N^2) — 퀵 셀렉트
- 공간: O(N)