수 N개(A1, A2, ..., Ai)가 주어진다. A를 오름차순으로 정렬했을 때 앞에서부터 K번째에 있는 수를 구하는 프로그램을 작성하시오.
입력
1번째 줄에 N(1 ⪯ N ⪯ 5,000,000)과 K(1 ⪯ K ⪯ N), 2번째 줄에 A1, A2, ..., Ai이 주어진다 (-109 ⪯ Ai ⪯ 109).
출력
A를 정렬했을 때 앞에서부터 K번째에 있는 수를 출력한다.
예제 입력
5 2 // 데이터 개수, K번째 수
4 1 2 3 5
예제 출력
2
1단계 - 문제 분석하기
최대 범위가 5,000,000이므로 O(nlogn)의 시간 복잡도로 정렬을 수행하면 된다. 퀵 배열을 구현해 주어진 수를 오름차순 배열하고 K번째 수를 출력해보자. 단, 이 문제는 시간 복잡도가 민감하므로 어떤 값을 pivot으로 정하면 K번째 수를 더 빨리 구할 수 있을지 생각해보자.
pivot을 정하는 방법
• pivot == K : K번째 수를 찾은 것이므로 알고리즘을 종료한다.
• pivot > K : pivot의 왼쪽 부분에 K가 있으므로 왼쪽(S ~ pivot-1)만 정렬을 수행한다.
• pivot < K : pivot의 오른쪽 부분에 K가 있으므로 오른쪽(pivot+1 ~ E)만 정렬을 수행한다.
데이터가 대부분 정렬되어 있는 경우 앞쪽에 있는 수를 pivot으로 선택하면 불필요한 연산이 많아진다. 따라서 이 문제는 배열의 중간 위치를 pivot으로 설정하였다.
2단계 - 손으로 풀어 보기
1 중간 위치를 pivot으로 설정한 다음 맨 앞에 있는 값과 swap한다. pivot을 맨 앞으로 옮기는 이유는 i, j 이동을 편하게 하기 위함이다.
이어서 i와 j를 pivot을 제외한 그룹에서 왼쪽, 오른쪽 끝으로 정한다.
2 우선 j를 이동한다. j가 pivot보다 크면 j--연산을 반복한다. 그 결과 j는 1에 위치하게 된다. j를 이동한 후에는 i가 pivot보다 작으면서 i보다 j가 크면 i++ 연산을 반복한다. 현재의 경우 i와 j값이 같으므로 i는 이동하지 않는다.
3 pivot을 두 집합을 나눠 주는 위치, 즉 i와 j가 만난 위치와 swap한다.
4 K는 2이므로 더이상 정렬하지 않고 A[2]를 출력한다.
3단계 - sudo코드 작성하기
N(숫자의 개수) K(K번째 수)
A(숫자 데이터 저장 배열)
for(N만큼 반복) {
A배열 저장
}
퀵 소트 실행
K번째 데이터 출력
[별도 함수 구현 부분]
퀵 소트 함수(시작, 종료, K)
{
피벗 구하기 함수(시작, 종료)
if(피벗 == K)
종료
else if(피벗 > K)
퀵 소트 실행(시작, 피벗-1, K)
else
퀵 소트 실행(피벗+1, 종료, K)
}
피벗 구하기 함수(시작, 종료)
{
데이터가 2개인 경우 바로 비교하여 정렬
N(중앙값)
중앙값을 시작 위치와 swap
pivot을 시작 위치 값 A[S]로 저장
i(시작점), j(종료점)
while(i <= j){
피벗보다 큰 수가 나올 때까지 i++
피벗보다 작은 수가 나올 때까지 j--
찾은 i와 j를 swap
}
피벗 데이터를 나눠진 두 그룹의 경계 index에 저장하기
경계 index 리턴
}
4단계 - 코드 구현하기
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;
public class Q19 {
public static void main(String[] args) throws IOException{
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());
st = new StringTokenizer(br.readLine());
int[] A = new int[N];
for(int i = 0; i < N; i++){
A[i] = Integer.parseInt(st.nextToken());
}
quickSort(A,0,N-1,K-1);
System.out.println(A[K-1]);
}
public static void quickSort(int[] A, int S, int E, int K){
if(S < E){
int pivot = partition(A, S, E);
if(pivot == K)
return;
else if(pivot > K)
quickSort(A, S, pivot-1, K);
else
quickSort(A, pivot+1, E, K);
}
}
public static int partition(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.length-1)
i++;
if(i <= j)
swap(A, i++, j--);
}
A[S] = A[j];
A[j] = pivot;
return j;
}
public static void swap(int[] A, int i, int j){
int tmp = A[i];
A[i] = A[j];
A[j] = tmp;
}
}
- Do it! 알고리즘 코딩테스트 자바 편
