[24060] 알고리즘 수업 - 병합 정렬 1 (JAVA)

# 문제 설명

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

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

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

merge_sort(A[p..r]) { # A[p..r]을 오름차순 정렬한다.
    if (p < r) then {
        q <- ⌊(p + r) / 2⌋;       # q는 p, r의 중간 지점
        merge_sort(A, p, q);      # 전반부 정렬
        merge_sort(A, q + 1, r);  # 후반부 정렬
        merge(A, p, q, r);        # 병합
    }
}

# A[p..q]와 A[q+1..r]을 병합하여 A[p..r]을 오름차순 정렬된 상태로 만든다.
# A[p..q]와 A[q+1..r]은 이미 오름차순으로 정렬되어 있다.
merge(A[], p, q, r) {
    i <- p; j <- q + 1; t <- 1;
    while (i ≤ q and j ≤ r) {
        if (A[i] ≤ A[j])
        then tmp[t++] <- A[i++]; # tmp[t] <- A[i]; t++; i++;
        else tmp[t++] <- A[j++]; # tmp[t] <- A[j]; t++; j++;
    }
    while (i ≤ q)  # 왼쪽 배열 부분이 남은 경우
        tmp[t++] <- A[i++];
    while (j ≤ r)  # 오른쪽 배열 부분이 남은 경우
        tmp[t++] <- A[j++];
    i <- p; t <- 1;
    while (i ≤ r)  # 결과를 A[p..r]에 저장
        A[i++] <- tmp[t++];
}

---

입력

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

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

---

출력

배열 A에 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());
        }

        mergeSort(A, 0, N-1);

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

    public static void mergeSort(int[] arr, int start, int end) {
        if (start < end) {
            int middle = (start + end) / 2;
            mergeSort(arr, start, middle);
            mergeSort(arr, middle+1, end);
            merge(arr, start, middle, end);
        }
    }

    public static void merge(int[] arr, int start, int middle, int end) {
        int[] tmp = new int[end-start+1];
        int p1 = start;
        int p2 = middle+1;
        int tmp_p = 0;

        while (p1 <= middle && p2 <= end) {
            if (arr[p1] <= arr[p2]) {
                tmp[tmp_p++] = arr[p1++];
            }
            else {
                tmp[tmp_p++] = arr[p2++];
            }
        }

        while (p1 <= middle) {
            tmp[tmp_p++] = arr[p1++];
        }

        while (p2 <= end) {
            tmp[tmp_p++] = arr[p2++];
        }

        p1 = start;
        tmp_p = 0;

        while (p1 <= end) {
            cnt++;
            if (cnt == K) {
                System.out.print(tmp[tmp_p]);
            }

            arr[p1++] = tmp[tmp_p++];
        }
    }
}

 

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

start가 end보다 작으면 집합의 중간 위치 middle을 (start + end) / 2로 설정한다.

middle을 기준으로 집합을 나눠 병합 정렬을 한다. (start~middle, middle+1~end)

 

merge 함수를 이용하여 두 집합을 정렬하며 합친다.

정렬할 배열, 시작 위치, 중간 위치, 끝 위치를 받는다.

두 집합을 합친 데이터를 저장할 tmp를 만든다.

각 집합의 포인터로 사용할 p1, p2와 tmp의 포인터로 사용할 tmp_p를 만든다.

 

p1이 middle보다 작거나 같고 p2가 end보다 작거나 같을 때까지 p1, p2가 가리키는 값 중 더 작은 값을 tmp에 넣는다.

while문이 끝나면 각 집합에 남은 값들을 tmp에 추가해준다.

 

두 집합의 병합이 끝나면 tmp에 있는 데이터를 기존 배열 A로 옮긴다.

이때 배열에 값을 저장하는 횟수 cnt를 1 증가한다.

만약 cnt가 K와 같으면, 현재 저장하는 숫자를 출력한다.

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