[백준] 11004 K 번째 수 quick selection sorting으로 풀기

1. 문제 설명

문제 링크

2. 접근 방식

quick 정렬을 이용해 문제를 풀었다. 하지만 기본적인 quick sort로 문제를 풀면 시간초과가 난다. quick sort의 최악 시간복잡도는 O(n^2) 인데, 데이터의 크기가 5,000,000 이라 n^2를 하면 2 억번의 계산을 훌쩍 넘기 때문이다.
따라서 O(nlogn)의 다른 정렬을 쓰거나, quick sort를 최적화 해줘야 한다.

나는 quick sort를 최적화하는 방향을 선택했다. 해당 문제는 k번째 index의 수가 무엇인지만 구하면 되는 문제이므로, 모든 배열을 정렬시킬 필요가 없다. 따라서

• 이분 탐색처럼 구역을 나눠서, 필요 없는 부분의 정렬은 하지 않는다.
• K번째 수를 알아내면 거기서 종료한다.

ⓐ 현재 정렬하려는 범위의 중앙값을 pivot으로 설정한다.

ⓑ mid와 start의 값을 바꿔서 pivot 값이 배열 맨 앞으로 올 수 있도록 한다. (pivot의 적정 위치 찾는 계산 시, 혼선이 없기 위해)

ⓒ 이제 start +1 ~ end 범위를 움직여 가면서 pivot 값이 들어갈 적정 위치를 찾으면 된다.
이때 i = start + 1, j = end라고 하겠다.

(c-1) j부터 움직인다. j가 가르키는 위치의 값이 pivot값보다 크면 j--로 왼쪽으로 이동한다.

(c-2) 다음은 i를 움직인다. i가 가르키는 값이 pivot 값보다 작으면 i++로 오른쪽으로 이동한다. (해당 예제는 i가 더 이상 움직일 곳이 없다.)

(c-3) 만약 j나 i가 움직이던 도중, j가 가르키는 값이 pivot 보다 작을 경우, 혹은 i가 가르키는 값이 pivot보다 클 경우 멈춰 있을 것이다. i < j 인 상황에서 i와 j가 멈춰있을 경우, 두 값의 위치를 바꾸고, i++, j--로 이동한다. (pivot 기준으로 그보다 큰 값은 오른쪽으로, 그보다 작은 값은 왼쪽으로 밀기)

ⓓ arr[start]와 arr[j]의 값을 swap 한다.
왜 i가 아니고 j인가? 왜냐하면, j의 경우 종국에 지칭하는 값은 pivot보다 작거나 같은 값이다. 왜냐하면 pivot보다 큰 값을 보면 언제나 왼쪽으로 자리 이동을 했기 때문이다. 그래서 j와 start의 값을 자리 바꿈 해도, pivot 기준 정렬이라는 quick-sort의 의미는 벗어나지 않는다.

ⓔ pivot의 index와 K를 비교한다.
현재 pivot만은 정렬이 된 상태이다.
pivot의 index를 P라고할 때, P==K이면 K번째 수를 찾았음으로 정렬을 종료하고 출력한다.
P < K 이면 pivot의 위치 오른편에 K가 존재하기 때문에 그 부분만 정렬을 이어간다.
P > K 이면 pivot의 위치 왼편에 K가 존재하기 떄문에 그 부분만 정렬한다.
이때 재귀를 이용한다.

3. 코드 분석

import java.io.*;
import java.util.StringTokenizer;

public class Main {
    public static void main(String[] args) throws IOException {
        // 0. 값 입력
        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()) - 1;
        int [] arr = new int [N];
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }

        quickSort(arr, 0, N-1, K);
        System.out.println(arr[K]);

    }

    public static void quickSort(int [] arr, int start, int end, int K) {
        if(start < end) {
            int pivot = partition(arr, start, end);

            if(pivot == K) return;
            else if (K < pivot) quickSort(arr, start, pivot-1, K);
            else quickSort(arr,pivot+1, end, K);
        }
    }


    // partition: 정해진 범위 내의 pivot 값이 들어갈 위치를 알려 주는 함수
    //            pivot 값을 기준으로 좌 우로 영역이 갈라짐.
    public static int partition(int[] arr, int start, int end) {
        // 만약 범위가 1이면, 둘 간의 정렬만 하면 됨.
        if(start+1==end){
            if(arr[start] > arr[end]) swap(arr, start, end);
            return end;
        }

        int mid = (start+end)/2; // pivot의 index
        swap(arr,start,mid); //    pivot 값을 왼쪽 끝으로 밀어서, 밑에서 이루어질, pivot의 진짜 위치 찾기 계산을 헷갈리지 않게 만듦
        int pivot = arr[start];

        int i = start+1; int j = end;

        while (i <= j) { // 중요! J가 완전히 pivot보다 큰 영역을 빠져 나와야 하므로, i == j 를 넘어서서 서로 어긋나도록 해야한다.
            while (j>= start+1 && pivot < arr[j]){ // 오른쪽 값이 기준 값보다 크면 내버려 두고 다음으로 넘어감
                j--;
            }

            while (i<=end&&pivot>arr[i]) { // 왼쪽 값이 기준 값보다 작을 경우, 그대로 두고 다음으로 넘어간다.
                i++;
            }
            if(i <= j){ // 둘 다 멈춰 있는데, 아직 서로 안 만난 경우, 왼쪽 값은 기준 값보다 크고, 오른쪽 값은 기준 값보다 작은 것임. 따라서 서로 swap
                swap(arr,i++,j--);
            }
        }
        arr[start] = arr[j];
        arr[j] = pivot;
        return j;
    }

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

4. 성장하기

이런 식으로 정렬이 필요한 부분만 선택해서 정렬하는 것을 quick-selection 정렬이라고 한다.