[24090] 알고리즘 수업 - 퀵 정렬 1 (JAVA)

# 문제 설명

오늘도 서준이는 퀵 정렬 수업 조교를 하고 있다. 아빠가 수업한 내용을 학생들이 잘 이해했는지 문제를 통해서 확인해보자.

N개의 서로 다른 양의 정수가 저장된 배열 A가 있다. 퀵 정렬로 배열 A를 오름차순 정렬할 경우 배열 A에 K 번째 교환되는 수를 구해서 우리 서준이를 도와주자.

크기가 N인 배열에 대한 퀵 정렬 의사 코드는 다음과 같다.

quick_sort(A[p..r]) { # A[p..r]을 오름차순 정렬한다.
    if (p < r) then {
        q <- partition(A, p, r);  # 분할
        quick_sort(A, p, q - 1);  # 왼쪽 부분 배열 정렬
        quick_sort(A, q + 1, r);  # 오른쪽 부분 배열 정렬
    }
}

partition(A[], p, r) {
    x <- A[r];    # 기준원소
    i <- p - 1;   # i는 x보다 작거나 작은 원소들의 끝지점
    for j <- p to r - 1  # j는 아직 정해지지 않은 원소들의 시작 지점
        if (A[j] ≤ x) then A[++i] <-> A[j]; # i값 증가 후 A[i] <-> A[j] 교환
    if (i + 1 != r) then A[i + 1] <-> A[r]; # i + 1과 r이 서로 다르면 A[i + 1]과 A[r]을 교환
    return i + 1;
}

---

입력

첫째 줄에 배열 A의 크기 N(5 ≤ N ≤ 10,000), 교환 횟수 K(1 ≤ K ≤ 108)가 주어진다.

다음 줄에 서로 다른 배열 A의 원소 A1, A2, ..., AN이 주어진다. (1 ≤ Ai ≤ 109)

---

출력

K 번째 교환되는 두 개의 수를 작은 수부터 한 줄에 출력한다. 교환 횟수가 K 보다 작으면 -1을 출력한다.

 

# 정답 코드

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

public class Main {
    static int cnt = 0;
    static int K = 0;

    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()); //배열의 크기
        K = Integer.parseInt(st.nextToken()); //교환 횟수

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

        quickSort(A, 0, N-1);

        if (cnt < K) {
            System.out.print("-1");
        }
    }

    public static void quickSort(int[] arr, int start, int end) {
        if (start < end) {
            int pivot = partition(arr, start, end);
            quickSort(arr, start, pivot-1); //왼쪽 분리 집합 정렬
            quickSort(arr, pivot+1, end); //오른쪽 분리 집합 정렬
        }
    }

    public static int partition(int[] arr, int start, int end) {
        int pivot = arr[end];
        int idx = start-1; //pivot보다 작은 데이터들의 끝지점

        for (int i = start; i < end; i++) {
            if (arr[i] <= pivot) {
                idx++;
                swap(arr, idx, i);
            }
        }

        if (idx+1 != end) {
            swap(arr, idx+1, end);
        }

        return idx+1;
    }

    public static void swap(int[] arr, int a, int b) {
        int num = arr[a];
        arr[a] = arr[b];
        arr[b] = num;

        cnt++;
        if (cnt == K) {
            System.out.print(arr[a] + " " + arr[b]);
        }
    }
}

 

quickSort 함수를 만들어 정렬할 배열, 시작 위치, 끝 위치를 입력하고 퀵 정렬을 한다.

start가 end보다 작으면 partition 함수로 pivot 값의 위치를 받는다.

pivot을 기준으로 분리된 왼쪽 분리 집합과 오른쪽 분리 집합도 퀵 정렬을 한다.

 

partition 함수는 정렬할 배열, 시작 위치, 끝 위치를 입력받아 정렬을 한다.

pivot은 집합에서 가장 오른쪽 값으로 정하므로 arr[end]로 설정한다.

pivot보다 작은 데이터들의 끝지점을 저장할 idx를 만든다.

 

i를 start부터 end-1까지 1씩 증가하게 하여 데이터를 순회한다.

만약 현재 i가 가리키는 값이 pivot보다 작거나 같으면 idx를 1 증가하고 idx가 가리키는 값과 swap한다.

 

pivot보다 작은 값들을 앞으로 이동시키는 작업이 끝나면, pivot을 적절한 위치에 삽입하면 된다.

• 만약 모든 값이 pivot보다 작으면(idx+1 == end), pivot을 이동시킬 필요가 없다.
• idx+1 != end이면, pivot보다 작은 데이터와 큰 데이터 사이(idx+1)에 pivot을 삽입해야 한다.
  idx+1 위치에 있는 값을 pivot이 있는 위치(end)와 swap한다.

pivot 이동이 끝나면, 현재 위치(idx+1)을 리턴한다.

 

교환은 swap에서만 일어나므로 swap이 일어날 때마다 cnt를 1 증가한다.

만약 cnt가 K와 같으면, 현재 교환한 값을 작은 순서대로 출력한다.

퀵 정렬이 모두 끝나고 cnt가 K보다 작으면, -1을 출력한다.