[Python] 병합 정렬, 재귀 - [백준 24060]

💡 문제 요약

🥈 난이도 실버3

• N개의 서로 다른 양의 정수가 저장된 배열 A가 있다.
• 병합정렬로 배열 A를 오름차순 정렬할 경우 배열 A에 K번째 저장되는 수를 구한다.
• 배열의 길이가 K보다 작으면 -1을 출력한다.

병합정렬 의사코드

입력/출력

📌 병합 정렬(merge sort)이란?

병합 정렬은 분할 정복 기법과 재귀 알고리즘을 이용한 정렬 알고리즘이다.
주어진 배열을 원소가 하나 밖에 남지 않을 때까지 계속 둘로 쪼갠 후 다시 크기 순으로 재배열 하며 정렬한 후 원래의 하나의 배열로 병합하는 방법이다.

예시) 주어진 배열이 아래와 같다.
[6, 4, 1, 3, 8, 2, 5, 7]

1. 하나의 배열을 둘로 쪼갠다
   [6, 4, 1, 3] [8, 2, 5, 7]

2. 배열의 길이가 1이 될 때까지 다시 둘로 쪼갠다.
   [6, 4] [1, 3] [8, 2] [5, 7]
[6] [4] [1] [3] [8] [2] [5] [7]

3. 더 이상 쪼갤 수 없어지면 두 개씩 크기를 비교하며 병합한다.
   [4, 6] [1, 3] [2, 8] [5, 7]
[1, 3, 4, 6] [2, 5, 7, 8]
[1, 2, 3, 4, 5, 6, 7, 8]

👩🏻‍💻 문제 해결

해당 문제는 단순히 정렬된 배열만이 아닌 완성된 배열의 k번째로 저장된 원소를 구하는 문제다.
그래서 마지막에 출력 시 K번째의 원소를 출력해주면 된다!

의사코드에서는 merge부분과 sort부분을 따로 썼지만 한번에 쓰는 것도 가능하다.

import sys

def mergeSort(lst):
    if len(lst) == 1:
        return lst

    mid = (len(lst) + 1) // 2
    # 왼쪽, 오른쪽 리스트 각각 정렬 - 재귀
    left = mergeSort(lst[:mid])
    right = mergeSort(lst[mid:])

    lst2 = []
    i, j = 0, 0

    # 병합
    while i < len(left) and j < len(right):
        # 오름차순 정렬이므로 작은 값을 저장
        if left[i] < right[j]:
            lst2.append(left[i])
            ans.append(left[i])
            i += 1
        else:
            lst2.append(right[j])
            ans.append(right[j])
            j += 1
    # 왼쪽 배열이 남은 경우
    while i < len(left):
        lst2.append(left[i])
        ans.append(left[i])
        i += 1
    # 오른쪽 배열이 남은 경우
    while j < len(right):
        lst2.append(right[j])
        ans.append(right[j])
        j += 1

    return lst2


N, K = map(int, sys.stdin.readline().split())
A = list(map(int, sys.stdin.readline().split()))

ans = []
mergeSort(A)

if len(ans) >= K:
    print(ans[K - 1])
else:
    print(-1)

위의 방식은 인덱스를 활용하였지만 또 다른 방법으로는 cnt변수를 사용해 저장할 때 마다 카운트를 1씩 증가시키고 이 cnt값이 K와 같아지면 그 때의 원소를 바로 출력하는 방법도 있다!

의사코드가 주어져서 어려운 문제는 아니었지만 정렬 알고리즘을 활용한 문제였기에 간단하게 포스팅을 해보았다:)
끝!