[백준/Python] 24060 알고리즘 수업 - 병합 정렬 1

[백준/Python] 24060 알고리즘 수업 - 병합 정렬 1

def merge_sort(arr, p, r, k, storage):
    if p < r:
        q = (p + r) // 2
        merge_sort(arr, p, q, k, storage)
        merge_sort(arr, q + 1, r, k, storage)
        merge(arr, p, q, r, k, storage)

def merge(arr, p, q, r, k, storage):
    temp = []
    i, j = p, q + 1
    while i <= q and j <= r:
        if arr[i] <= arr[j]:
            temp.append(arr[i])
            i += 1
        else:
            temp.append(arr[j])
            j += 1
        storage['count'] += 1
        if storage['count'] == k:
            storage['kth_value'] = temp[-1]
    while i <= q:
        temp.append(arr[i])
        i += 1
        storage['count'] += 1
        if storage['count'] == k:
            storage['kth_value'] = temp[-1]
    while j <= r:
        temp.append(arr[j])
        j += 1
        storage['count'] += 1
        if storage['count'] == k:
            storage['kth_value'] = temp[-1]
    for i, val in enumerate(temp):
        arr[p + i] = val

N, K = map(int, input().split())
A = list(map(int, input().split()))

storage = {'count': 0, 'kth_value': -1}

merge_sort(A, 0, N-1, K, storage)

print(storage['kth_value'])

1. merge_sort 함수는 배열 arr의 부분 배열 arr[p..r]을 병합 정렬 방식으로 정렬한다. 배열은 0-인덱스 기준으로 접근한다.
2. merge 함수는 정렬된 두 부분 배열 arr[p..q]와 arr[q+1..r]을 하나의 정렬된 배열로 병합하는 동시에, 요소가 저장될 때마다 storage['count']를 1씩 증가한다.
3. 만약 storage['count']가 입력으로 주어진 K와 같아지면, 그 시점에 저장되는 수를 storage['kth_value']에 저장한다.
4. 병합 정렬이 완료된 후, storage['kth_value']에는 K번째로 저장된 값이 저장되어 있다. 만약 K번째 저장이 발생하지 않았다면, -1이 출력된다.