병합정렬 (Merge-Sort) with Python

📖 문제 설명

문제 : https://www.acmicpc.net/problem/24060

• 오늘도 서준이는 병합 정렬 수업 조교를 하고 있다. 아빠가 수업한 내용을 학생들이 잘 이해했는지 문제를 통해서 확인해보자.
• 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++]; 
}

병합정렬이란

• 리스트를 반복적으로 나누고, 각 부분을 정렬한 후 다시 병합하여 전체 리스트를 정렬합니다.
• 병합 정렬은 안정적인 정렬 방법 중 하나로, 최악의 경우에도 시간 복잡도가 O(n log n)인 것이 특징입니다.

코드

import sys

n, k = map(int, sys.stdin.readline().split())
arr = list(map(int, sys.stdin.readline().split()))

def merge_sort(arr, start, end): # 리스트를 잘게 쪼개기 
    if start < end:
        mid = (start + end) // 2
        merge_sort(arr, start, mid)
        merge_sort(arr, mid + 1, end)
        merge(arr, start, mid, end)

def merge(arr, start, mid, end): 
    tmp = [0] * (end - start + 1) # 정렬된 배열을 담기 위한 배열 
    first_arr_index = start # 첫번째 조개진 배열 스타트 지점 
    second_arr_index = mid + 1 # 두번째 조개진 배열 스타트 지점
    merge_arr_index = 0 # 정렬된 배열을 담기위한 인덱스 

    while first_arr_index <= mid and second_arr_index <= end: #조개진 배열을 넘어가지 않도록 예외 처리 
        if arr[first_arr_index] <= arr[second_arr_index]: # 크기 비교
            tmp[merge_arr_index] = arr[first_arr_index] # 배열에 순차적으로 삽입
            first_arr_index += 1 # 다음 인덱스 번호로 증가 

        else:
            tmp[merge_arr_index] = arr[second_arr_index]
            second_arr_index += 1

        merge_arr_index += 1 # 다음 인덱스에 정렬된 값 삽입하기 위한 증가 

    while first_arr_index <= mid: # 만약 첫번째 배열의 값이 아직 tmp 배열에 저장되지 않았을 경우 
        tmp[merge_arr_index] = arr[first_arr_index]
        first_arr_index += 1 
        merge_arr_index += 1

    while second_arr_index <= end: 
        tmp[merge_arr_index] = arr[second_arr_index]
        second_arr_index += 1
        merge_arr_index += 1

    for i in range(start, end + 1): # tmp배열을 arr 배열에 삽입 
        arr[i] = tmp[i - start]

merge_sort(arr, 0, len(arr) - 1) 

🧑‍💻 코드 풀이

import sys

n, k = map(int, sys.stdin.readline().split())
arr = list(map(int, sys.stdin.readline().split()))

def merge_sort(arr, start, end):
    if start < end:
        mid = (start + end) // 2
        left_sort = merge_sort(arr, start, mid)
        if left_sort != -1:
            return left_sort

        right_sort = merge_sort(arr, mid + 1, end)
        if right_sort != -1:
            return right_sort

        return merge(arr, start, mid, end)

    return -1

def merge(arr, start, mid, end):
    global k
    tmp = [0] * (end - start + 1)
    first_arr_index = start
    second_arr_index = mid + 1
    merge_arr_index = 0

    while first_arr_index <= mid and second_arr_index <= end:
        if arr[first_arr_index] <= arr[second_arr_index]:
            tmp[merge_arr_index] = arr[first_arr_index]
            first_arr_index += 1

        else:
            tmp[merge_arr_index] = arr[second_arr_index]
            second_arr_index += 1

        merge_arr_index += 1
        k -= 1
        if k == 0:
            return tmp[merge_arr_index - 1]

    while first_arr_index <= mid:
        tmp[merge_arr_index] = arr[first_arr_index]
        first_arr_index += 1
        merge_arr_index += 1
        k -= 1
        if k == 0:
            return tmp[merge_arr_index - 1]

    while second_arr_index <= end:
        tmp[merge_arr_index] = arr[second_arr_index]
        second_arr_index += 1
        merge_arr_index += 1
        k -= 1
        if k == 0:
            return tmp[merge_arr_index - 1]

    for i in range(start, end + 1):
        arr[i] = tmp[i - start]

    return -1

result = merge_sort(arr, 0, len(arr) - 1)
print(result)

• 병합정렬과 코드와 크게 다르지 않습니다.
• 정렬된 배열을 저장하는 시점에 k -= 1 을 통해 값을 카운팅 해줍니다.
• 만약 k == 0 이면 tmp[merge_arr_index - 1] 리턴해줍니다.
• merge_sort 클래스에서 DFS 돌다가 리턴값이 -1 이 아닌값이 나오게되면 merge_sort는 종료됩니다.