문제 설명

문제 바로가기

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

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);
    }
}

분석

먼저, 의사 코드를 그대로 구현하여 풀 수 있겠다.

다른 풀이가 있을까 생각해보았지만 생각이 나지 않았다.

의사 코드를 구현하되 교환(swap) 횟수를 저장하고 교환 횟수가 K과 같으면 예외를 던져 처리하도록 하자.

생각해보기

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

public class Main {
    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());
        int K = Integer.parseInt(st.nextToken());
        int[] arr = new int[N+1];
        StringTokenizer arrSt = new StringTokenizer(br.readLine());
        for (int i = 1; i <= N; i++) arr[i] = Integer.parseInt(arrSt.nextToken());
        System.out.println(new Solution(arr, K).solution());
    }

    static class Solution {
        int count = 0;
        int K;

        int[] arr;

        public Solution(int[] arr, int K) {
            this.arr = arr;
            this.K = K;
        }

        public String solution() {
            try {
                heap_sort(arr);
            } catch (Exception e) {
                return toString();
            }
            return "-1";
        }

        private void heap_sort(int[] arr) {
            int n = arr.length - 1;
            buildMinHeap(arr, n);
            for(int i = n; i >= 2; i--) {
                swap(arr, 1, i);
                heapify(arr, 1, i - 1);
            }
        }

        private void buildMinHeap(int[] arr, int n) {
            for(int i = n/2; i >= 1; i--) heapify(arr, i, n);
        }

        private void heapify(int[] arr, int k, int n) {
            int left = 2 * k;
            int right = left + 1;
            int smaller = 0;
            if(right <= n) {
                if(arr[left] < arr[right]) smaller = left;
                else smaller = right;
            } else if (left <= n) {
                smaller = left;
            } else {
                return;
            }
            if(arr[smaller] < arr[k]) {
                swap(arr, smaller, k);
                heapify(arr, smaller, n);
            }
        }

        private void swap(int[] arr, int i, int j) {
            int t = arr[i];
            arr[i] = arr[j];
            arr[j] = t;
            if(++count == K) throw new RuntimeException();
        }

        @Override
        public String toString() {
            StringJoiner sj = new StringJoiner(" ");
            for (int i = 1; i < arr.length; i++) sj.add(Integer.toString(arr[i]));
            return sj.toString();
        }
    }
}

swap 함수 안에서 count를 증가시키고 조건에 맞으면 예외를 발생시킨다.
solution 함수에서 예외를 캐치하면 K번 swap한 상태의 배열을 문자열로 반환한다.
solution 함수에서 아무런 예외를 캐치하지 않으면 "-1"을 반환한다.

마무리

뭔가... 뭔가 좀 더 좋은 풀이가 있을까?
다름 사람들의 풀이를 좀 보려고 했더니 풀이가 별로 없고 다 나와 같이 의사코드를 그대로 구현한 풀이였다. 신생 문제인가 보다.
나중에 기억나면 다시 한번 뒤져봐야겠다.