백준 11004 - K번째 수

문제

시간 제한 2초

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

입력

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

둘째에는 $A_1, A_2, ..., A_N$이 주어진다. ($-10^9$ ≤ $A_i$ ≤ $10^9$)

출력

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

예제 입력 1

5 2
4 1 2 3 5

예제 출력 1

2

문제 풀이

먼저 시간복잡도를 살펴보면 N의 최댓값은 5,000,000 이다 정렬을 사용한다면 $O(nlogn)$ 이기 때문이다. Arrays.sort() 함수를 사용하여도 문제는 없을것 같다.

정렬한 후에는 K 번째 수를 구하면 되기 때문에 이를 코드에 적용하면 다음과 같다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

class Main {

    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());
        int[] A = new int[N];

        st = new StringTokenizer(br.readLine());
        for(int i = 0; i < N; i++){
            A[i] = Integer.parseInt(st.nextToken());
        }

        Arrays.sort(A);

        System.out.println(A[K - 1]);


    }

}

사실 책(Do it! 알고리즘 코딩 테스트 자바) 에서는 퀵 정렬을 이용하여 문제를 풀었지만 왜인지 모르게 시간초과가 계속 발생한다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class 백준K번째수 {

    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());
        int[] A = new int[N];
        st = new StringTokenizer(br.readLine());
        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]);

    }

    private static void quickSort(int[] A, int start, int end, int K) {
        if(start < end) {
            int pivot = partition(A, start, end);
            if(pivot == K)
                return;
            else if(K < pivot)
                quickSort(A, start, pivot - 1, K);
            else
                quickSort(A, pivot + 1, end, K);
        }
    }

    private static int partition(int[] A, int start, int end) {
        int middle = (start + end) / 2;
        swap(A, start, middle);             // 중앙값을 1번째 요소로 이동하기
        int pivot = A[start];
        int i = start, j = end;
        while(i < j){
            while(pivot < A[j]){            // 피벗보다 작은 수가 나올 때까지 j--
                j--;
            }
            while(i < j && pivot >= A[i]) { // 피벗보다 큰 수가 나올 때까지 i++
                i++;
            }
            swap(A, i, j);                  // 찾은 i와 j를 교환하기
        }
        // i == j 피벗의 값을 양쪽으로 분리한 가운데에 오도록 설정하기
        A[start] = A[i];
        A[i] = pivot;
        return i;
    }

    private static void swap(int[] A, int i, int j) {
        int temp = A[i];
        A[i] = A[j];
        A[j] = temp;
    }

}

회고

글을 언제 수정할지는 모르겠지만 분명 퀵 정렬로 푸는게 더 옳은 코드일 것 같다.
퀵 정렬 알고리즘을 배운다고 생각하여야 겠다.