[백준 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++]; 
}

 

입력

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

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

출력

배열 A에 K번째 저장 되는 수를 출력한다. 저장 횟수가 K 보다 작으면 -1을 출력한다.

 

예제 입력

입력 1
5 7
4 5 1 3 2

입력 2
5 13
4 5 1 3 2

4 5 1 3 2 -> 4 5 1 3 2 -> 4 5 1 3 2 -> 1 5 1 3 2 -> 1 4 1 3 2 -> 1 4 5 3 2 -> 1 4 5 2 2 -> 1 4 5 2 3 -> 1 4 5 2 3 -> 1 2 5 2 3 -> 1 2 3 2 3 -> 1 2 3 4 3 -> 1 2 3 4 5. 총 12회 저장이 발생하고 일곱 번째 저장되는 수는 3이다.

예제 출력

출력 1
3

출력 2
-1

 

 

---

 

풀이

병합 정렬이 어떤 정렬인지 먼저 설명하자면
4 5 1 3 2 배열이 주어지고,
배열 길이를 절반씩 쪼개서 길이를 1로 만들어 정렬 하는 방법입니다.

주어진 배열의 길이가 5이므로 절반은 2.5 이지만 소수점은 버립니다. (중간점 = 2)
중간점을 기준으로 왼쪽과 오른쪽으로 배열을 나눕니다.
4 5 / 1 3 2

다시 왼쪽 배열을 중간점을 기준으로 배열을 나눕니다. (왼쪽 배열 길이 2, 중간점 = 1)
4 / 5
길이가 1로 더이상 쪼개지지 않으니, 4 - 5를 비교해서 작은 수를 배열에 넣어줍니다.
[4, 5]

이번엔 오른쪽 배열을 중간점을 기준으로 배열을 나눕니다. (오른쪽 배열 길이 3, 중간점 = 1)
1 / 3 2
오른쪽 배열이 아직 길이가 1이 아니니까 한번 더 쪼개줍니다. (오른쪽 배열 길이 2, 중간점 = 1)
3 / 2
길이가 1로 더이상 쪼개지지 않으니 3 - 2를 비교해서 작은 수를 배열에 넣어줍니다.
그리고 왼쪽에 있던 길이가 1개인 숫자와 비교해서 배열에 넣어줍니다.
1 / [2, 3] → [1, 2, 3]

미지막으로 [4, 5] / [1, 2, 3] 을 비교해서 정렬합니다.
최종적으로 [1, 2, 3, 4, 5] 으로 정렬이 완료됩니다.

이렇게 해서 코드를 짜보면
1. 배열을 길이가 1까지 쪼갤 재귀함수 생성
    1-1. 길이가 1보다 작으면 배열 출력
    1-2. 중간지점을 저장하고, 중간점을 기준으로 왼쪽, 오른쪽으로 나누기
    1-3. 왼쪽, 오른쪽으로 나눈 배열을 다시 재귀적으로 호출해서 길이가 1보다 작을 때까지 쪼개기
    1-4. 왼쪽, 오른쪽으로 나눈 배열을 병합하는 함수에 대입
2. 쪼갠 배열을 정렬하여 병합하는 함수 생성
    2-1. 정렬 리스트를 선언
    2-2. 배열을 하나씩 늘려가서 비교할 인덱스 변수 선언
    2-3. 인덱스가 배열의 길이보다 작을 때까지 반복
    2-4. 왼쪽, 오른쪽 배열을 비교하여 작은 값을 정렬 리스트에 저장
    2.5. 인덱스 길이가 어떤 배열의 길이를 초과하는 경우 다른 한쪽 배열에 들어있는 값을 모두 정렬 리스트에 저장

병합 정렬 코드

def merge_sort(arr): # 배열을 쪼개는 함수
    if len(arr) <= 1:
        return arr
    
    mid = len(arr) // 2
    left = arr[:mid] # 중간 지점을 기준으로 왼쪽부분 인덱스 슬라이싱
    right = arr[mid:] # 중간 지점을 기준으로 오른쪽부분 인덱스 슬라이싱
    
    left_arr = merge_sort(left) # 쪼갠 왼쪽 배열을 재귀적으로 호출하여 중간 지점을 기준 쪼갬
    right_arr = merge_sort(right) # 쪼갠 오른쪽 배열을 재귀적으로 호출하여 중간 지점을 기준 쪼갬
    return merge(left_arr, right_arr) # 쪼개진 왼쪽, 오른쪽 배열을 병합하는 함수에 대입
    
def merge(left_arr, right_arr): # 왼쪽, 오른쪽 배열을 병합하는 함수
    list_sort = [] # 정렬 리스트 선언
    i, j = 0, 0 # 0부터 ~ 확인 할 인덱스 선언
    
    while i < len(left_arr) and j < len(right_arr): # 인덱스 길이보다 배열 길이가 크다면
        if left_arr[i] < right[j]: # 왼쪽의 i번째, 오른쪽의 j번째와 비교해서 작은 걸 정렬리스트에 추가
            list_sort.append(left_arr[i])
            i += 1 # 왼쪽 배열이 추가 됐다면 i를 +1
        else:
            list_sort.append(right_arr[j])
            j += 1 # 오른쪽 배열이 추가 됐다면 j를 +1
            
    # 왼쪽, 오른쪽 배열 중 원소가 남아있는게 있다면 정렬리스트에 모두 추가
    while i < len(left_arr):
        list_sort.append(left_arr[i])
        i += 1
    while j < len(right_arr):
        list_sort.append(right_arr[j])
        j += 1

	return list_sort # 정렬 리스트를 리턴​

병합 정렬 코드는 이렇게 되는데요.
이 문제에서는 배열의 크기와 저장 횟수가 주어지고,
병합 정렬을 했을 때 주어진 저장 횟수만큼 했을 때 저장되는 수를 구해야 합니다.
주어진 저장횟수가 총 횟수를 초과하면 -1을 출력하네요.

예제를 보면
배열의 크기 : 5, 저장 횟수 : 7
배열 : 4 5 1 3 2
4 5 1 3 2 -> 4 5 1 3 2 -> 4 5 1 3 2 -> 1 5 1 3 2 -> 1 4 1 3 2 -> 1 4 5 3 2 -> 1 4 5 2 2 -> 1 4 5 2 3 -> 1 4 5 2 3 -> 1 2 5 2 3 -> 1 2 3 2 3 -> 1 2 3 4 3 -> 1 2 3 4 5.

예제에서는 배열의 크기가 5일 때 중간점을 3으로 해서 구했네요.
그럼 코드에서 mid = len(arr) // 2 이 부분을 mid = (len(arr) + 1) // 2 로 고쳐줘야 겠네요.

그리고 7번째 횟수에 저장되는 수를 구해야 하는데,
함수 밖에 또 다른 리스트를 선언해서 병합 정렬하는 함수의 정렬리스트와 함께 값을 추가해 주기만 하면 됩니다.
그럼 그 리스트 안에 [4, 5, 1, 4, 5, 2, 3, 1, 2, 3, 4, 5] 이렇게 순서대로 값이 들어옵니다.
이 중에 7번째의 값이 출력되게 해줍니다.
길이가 12를 넘어가면 -1이 출력되게 해줍니다.

---

병합 정렬 코드는 윗부분에 자세히 설명 되어있습니다.

작성 코드

def merge_sort(arr):
    if len(arr) <= 1:
        return arr
    
    mid = (len(arr)+1) // 2
    left = arr[:mid]
    right = arr[mid:]

    left_sort = merge_sort(left)
    right_sort = merge_sort(right)
    return merge(left_sort, right_sort)

def merge(left_sort, right_sort):
    i, j = 0, 0 
    list_sort = []

    while i < len(left_sort) and j < len(right_sort):
        if left_sort[i] < right_sort[j]:
            list_sort.append(left_sort[i])
            answer.append(left_sort[i]) ##### 이부분이 추가됨 #####
            i += 1
        else:
            list_sort.append(right_sort[j])
            answer.append(right_sort[j]) ##### 이부분이 추가됨 #####
            j += 1

    while i < len(left_sort):
        list_sort.append(left_sort[i])
        answer.append(left_sort[i]) ##### 이부분이 추가됨 #####
        i += 1
    while j < len(right_sort):
        list_sort.append(right_sort[j])
        answer.append(right_sort[j]) ##### 이부분이 추가됨 #####
        j += 1
    return list_sort
    
n, k = map(int, input().split())
A = list(map(int, input().split()))
answer = [] # 새로운 리스트 선언
merge_sort(A) # 배열을 병합정렬 함수 안에 대입

if len(answer) >= k: # 길이가 k보다 크거나 같으면
    print(answer[k-1]) # 인덱스는 0부터 시작하니 k-1번째를 출력
else:
    print(-1) # 길이가 k보다 작으면 -1 출력