알고리즘 설계 실습 - k번째 교환(heap sort)

---

문제

N개의 양의 정수가 저장된 배열 A가 있습니다. 아래 의사 코드 (pseudo-code)를 사용한 힙 정렬 로 배열 A를 정렬할 경우 K 번째 교환되는 수와 K 번 교환된 직후의 배열 A의 모든 원소를 출력 해 보세요.

N개의 양의 정수가 저장된 배열에 대한 힙 정렬 의사 코드는 아래와 같습니다.

heap_sort(A[1..n]) { # A[1..n]을 정렬한다. build_min_heap(A, n);
for i <- n downto 2 {
A[1] <-> A[i]; # 원소 교환
        heapify(A, 1, i - 1);
    }
}
build_min_heap(A[], n) { fori<- ⌊n/2⌋ downto1
        heapify(A, i, n)
}
# A[k]를 루트로 하는 트리를 최소 힙 성질을 만족하도록 수정한다.
# A[k]의 두 자식을 루트로 하는 서브 트리는 최소 힙 성질을 만족하고 있다. # n 은 배열 A 의 전체 크기이며 최대 인덱스를 나타낸다.
heapify(A[], k, n) {
left <- 2k; right <- 2k + 1;
    if (right ≤ n) then {
if (A[left] < A[right]) then smaller <- left; else smaller <- right;
}
else if (left ≤ n) then smaller <- left; else return;
# 최소 힙 성질을 만족하지 못하는 경우 재귀적으로 수정한다. if (A[smaller] < A[k]) then {
        A[k] <-> A[smaller];
heapify(A, smaller, n); }
}

입력

첫째 줄에 배열 A의 크기 N(5≤N≤5000), 교환 횟수 K(1≤K≤108)가 주어집니다. 다음 줄에 서로 다른 배열 A의 원소 A1,A2,...,AN이 주어집니다.(1≤Ai ≤108)

출력

총 교환 횟수가 K 보다 클 경우, 첫 번째 줄에 K 번째 교환되는 두 개의 수를 작은 수부터 공백 을 포함하여 출력하고, 두 번째 줄에 배열 A에 K 번 교환이 발생한 직후의 배열 A의 모든 원소를 공백을 포함하여 출력하세요.

총 교환 횟수가 K 보다 작으면 첫 번째 줄에만 -1을 출력하세요.

---

Solution

 

import sys

input = sys.stdin.readline  # 속도 개선

# 교환은 heapify과정과 최소값 추출 직후의 case에서 발생
# 변수를 공유하는 것이 코드 작성에 유리하여 heap_sort함수 내에 heapify함수와 build_를 정의
def heap_sort(A, N, K):
    # 주어진 노드 k에 대해 최소 힙 속성을 유지시키는 heapify함수
    def heapify(A, k, n, exchanges):
        smallest = k  # smallest인덱스를 초기화
        left = 2 * k  # 왼쪽 자식 노드의 인덱스
        right = 2 * k + 1  # 오른쪽 자식 노드의 인덱스

        # 왼쪽 자식 노드와 현재 노드의 값을 비교하여, 더 작은 값을 가진 노드의 인덱스를 저장
        if left <= n and A[left - 1] < A[smallest - 1]:
            smallest = left
        
        # 오른쪽 자식 노드와 현재 '가장 작은 값'을 가진 노드를 비교하여, 더 작은 값을 가진 노드의 인덱스를 저장
        if right <= n and A[right - 1] < A[smallest - 1]:
            smallest = right

        # 가장 작은 값을 가진 노드가 현재 노드가 아닌 경우, 두 노드의 값을 교환하고 교환 내역을 저장
        if smallest != k:
            A[k - 1], A[smallest - 1] = A[smallest - 1], A[k - 1]
            exchanges.append((min(A[k - 1], A[smallest - 1]), max(A[k - 1], A[smallest - 1])))
            
            # K번째 교환을 찾은 경우, True를 반환
            if len(exchanges) == K:
                return True
            
            # 하위 트리에 대해 재귀적으로 heapify 함수를 호출
            return heapify(A, smallest, n, exchanges)
        return False
    
    #주어진 배열 A를 최소 힙으로 구성하는 build_min_heap함수
    def build_min_heap(A, n, exchanges):
        # 가장 마지막 부모 노드부터 시작하여, 배열의 모든 노드에 대해 heapify 함수를 호출
        for i in range(n // 2, 0, -1):
            if heapify(A, i, n, exchanges):
                return True  # K번째 교환을 찾은 경우, True를 반환
        return False

    exchanges = []  # 교환 내역을 저장할 리스트를 초기화

    # 배열 A를 최소 힙으로 구성합 K번째 교환을 찾은 경우, 해당 교환 내역과 힙 상태를 출력
    if build_min_heap(A, N, exchanges):
        print(*exchanges[-1])
        print(*A)
        return exchanges[-1], A

    # 최소 힙에서, smallest를 맨 끝 노드와 교환하고, 힙의 크기를 줄인 후 다시 최소 힙을 구성
    for i in range(N, 1, -1):
        if A[0] != A[i - 1]:
            A[0], A[i - 1] = A[i - 1], A[0]
            exchanges.append((min(A[0], A[i - 1]), max(A[0], A[i - 1])))
            
            # K번째 교환을 찾은 경우, 해당 교환 내역과 힙 상태를 출력합니다.
            if len(exchanges) == K:
                print(*exchanges[-1])
                print(*A)
                return exchanges[-1], A
        if heapify(A, 1, i - 1, exchanges):
            print(*exchanges[-1])
            print(*A)
            return exchanges[-1], A

    # K번째 교환이 존재하지 않는 경우, -1을 출력
    print(-1)

# 배열 A와 정수 K를 입력받아 heap_sort 함수를 호출
N, K = map(int, input().split())
A = list(map(int, input().split()))
heap_sort(A, N, K)

그렇게 어려운 문제는 아닌거 같은데 정말 애를 많이 먹었다. 주어진 의사코드대로는 코드를 작성하지 않고 중첩함수를 이용하여 좀더 가독성을 높인 코드를 작성하였다. 주석을 매우 자세히 달아놓았기에 추가설명은 필요없을 것으로 보인다. 핵심 코드 진행은 초기의 최소 힙 구성 과정에서 swap 횟수를 카운팅하고 이후로는 최소값을 찾을 때 마다 진행 되는 swap과 heapify과정에서의 swap을 꼼꼼히 카운팅하다가 카운터가 k가 되면 중지하고 출력하는 것이다. 

힙 정렬에 대한 글은 이전에 정리해둔 내용을 첨부한다.

Algorithm Design - 셸 정렬 / 퀵 정렬 / 힙 정렬 / 계수 정렬 / 기수 정렬

 

---