[백준 실버5] 24174 알고리즘 수업 - 힙 정렬2 / Java

---

문제

문제 설명

오늘도 서준이는 최소 힙 기반 힙 정렬 수업 조교를 하고 있다. 아빠가 수업한 내용을 학생들이 잘 이해했는지 문제를 통해서 확인해보자.

N개의 서로 다른 양의 정수가 저장된 배열 A가 있다. 힙 정렬로 배열 A를 오름차순 정렬할 경우 배열 A의 원소가 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) {
    for i <- ⌊n / 2⌋ downto 1
        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 ≤ 500,000), 교환 횟수 K(1 ≤ K ≤ 10⁸)가 주어진다.

다음 줄에 서로 다른 배열 A의 원소 A₁, A₂, ..., A_N이 주어진다. (1 ≤ A_i ≤ 10⁹)

출력

배열 A에 K 번 교환이 발생한 직후의 배열 A를 한 줄에 출력한다. 교환 횟수가 K 보다 작으면 -1을 출력한다.

---

아이디어

단순한 array에서의 heap 정렬을 구현하는 문제인데, 아직 힙에 대해 이해도가 부족하여 별다른 아이디어는 없는 것같다.

---

소스코드

// 백준 24174번 힙 정렬

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

// max heap sort = 오름차순
// min heap sort = 내림차순

class Main{
    public static int cnt = 0, target = 0;
    public static boolean flag = false;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st =new StringTokenizer(br.readLine());

        int N = Integer.parseInt(st.nextToken());
        target = Integer.parseInt(st.nextToken());

        st = new StringTokenizer(br.readLine());
        int[] arr = new int[N];
        for (int i = 0; i < N ; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }
        heapSort(arr);
        if(!flag){
            System.out.println(-1);
        }
    }

    // swap
    // 해당 문제에서는 cnt에 맞춰서 StringBuilder로 출력
    public static void swap(int[] arr ,int i, int j){
        cnt++;
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;

        if(target == cnt){
            StringBuilder sb = new StringBuilder();
            for (int k = 0; k < arr.length; k++) {
                sb.append(arr[k]).append(" ");
            }
            System.out.println(sb);
            flag= true;
        }
    }

    // min heap 만드는 과정
    public static void min_heapify(int[] arr, int parentIdx, int lastIdx){
        int left= parentIdx * 2 + 1;
        int right = parentIdx * 2 + 2;
        int largestIdx = parentIdx;

        if(left <= lastIdx && arr[left] < arr[largestIdx]){
            largestIdx = left;
        }
        if(right <= lastIdx && arr[right] < arr[largestIdx]){
            largestIdx= right;
        }

        if(parentIdx != largestIdx){
            swap(arr,parentIdx,largestIdx);
            min_heapify(arr,largestIdx,lastIdx);
        }

    }
    // max heap 만드는 과정
    public static void max_heapify(int[] arr, int parentIdx, int lastIdx){
        int left= parentIdx * 2 + 1;
        int right = parentIdx * 2 + 2;
        int largestIdx = parentIdx;

        if(left <= lastIdx && arr[left] > arr[largestIdx]){
            largestIdx = left;
        }
        if(right <= lastIdx && arr[right] > arr[largestIdx]){
            largestIdx = right;
        }

        if(parentIdx != largestIdx){
            swap(arr,parentIdx,largestIdx);
            max_heapify(arr,largestIdx,lastIdx);
        }
    }

    // 부모노드 찾기
    public static int getParent(int idx){
        return (idx - 1) / 2;
    }
    public static void heapSort(int[] arr){
        int size = arr.length;

        if(size < 2) return;

        // 부모노드 찾기
        int parentIdx = getParent(size - 1);

        // heap 만드는 과정
        for (int i = parentIdx; i >= 0 ; i--) {
            // max_heapify(arr,i,size - 1);
            // 부모노드부터, 0번 인덱스까지
            // 자식 노드의 -1 level부터 0 level까지 탐색하면서, 최대 및 최소 힙 생성
            min_heapify(arr,i,size-1);

        }

        // 정렬하는 과정
        for (int i = size - 1; i > 0  ; i--) {
            // 맨 앞 노드와 뒤 노드 위치 전환
            swap(arr,0,i);
            // max_heapify(arr,0,i-1);
            min_heapify(arr,0,i-1);
        }
    }

}

---

결과

결과

---

후기

힙에 익숙치가 않아서 다시 한 번 공부해보면서 학습하는 기회가 되었던 것같다.

특히, 한 줄씩 주석을 적어가면서 힙에 대해 이해하고 구현할 수 있었던 문제인것같다.