[BOJ/JAVA] 24174. 알고리즘 수업 - 힙 정렬 2

문제 설명

문제 링크

접근 방법

뭔가 복잡하게 생각할 거는 없었다. 그냥 의사코드 대로 그대로 구현하면 되는데, 다만 고려해야 할 점은, count 가 K 에 도달했을 때의 처리이다. 처음에는 조금 편하게 하려고 배열을 LinkedList에 저장해서 Collections.swap()을 이용했고, 그렇게 하다보니 K에 도달해도 끝까지 실행하는 문제가 발생했다. 그 때문인지는 몰라도 시간 초과가 떠서, 정직하게 코드를 다시 고쳐서 배열을 이용하는 방법으로 구현했다. 그 과정에서 시간을 단축하기 위해 swap() 메소드에서 K 에 도달하면 출력 뒤 강제 종료하는 코드를 작성하였다.

소스 코드

import java.io.*;
import java.util.*;

public class Main {
    static int count = 0;
    static int K = 0;

    public static void swap(int[] arr, int a, int b) {
        int tmp = arr[a];
        arr[a] = arr[b];
        arr[b] = tmp;
        count++;

        if(K == count) {
            StringBuilder sb = new StringBuilder();
            for(int i=1; i<arr.length; i++) {
                sb.append(arr[i]).append(" ");
            }
            System.out.println(sb.toString());
            System.exit(0);
        }
    }

    public static void heapSort(int[] a) {
        int n = a.length-1;
        buildMinHeap(a, n);
        for (int i=n; i>=2; i--) {
            swap(a, 1, i);
            heapify(a, 1, i-1);
        }
    }

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

    public static void heapify(int[] a, int k, int n) {
        int left = 2*k;
        int right = 2*k + 1;
        int min;
        if (right <= n ) {
            if (a[left] < a[right]) min = left;
            else min = right;
        } else if (left <= n) min = left;
        else return;

        if(a[min] < a[k]) {
            swap(a, min, k);
            heapify(a, min, n);
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        //BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int n = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());
        st = new StringTokenizer(br.readLine());

        int[] a = new int[n+1];
        for(int i=1; i<n+1; i++) a[i] = Integer.parseInt(st.nextToken());

        heapSort(a);
        System.out.println("-1");

        br.close();
    }
}

제출 결과

마무리하며

문제 제목처럼, 대학생 때 알고리즘 수업에서 비슷한 걸 했던 기억이 떠올라서 재밌는 문제였다. (그때 좀 열심히 할걸)