[백준] 24174번 - 힙 정렬 2

💎 문제 링크

24174번: 알고리즘 수업 - 힙 정렬 2

 

📝 문제

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

 

✍🏻 입력 & 출력

 

🤔 풀기 전에 생각하기

슈도 코드가 제시되어 있기 때문에 그걸 자바 코드로 바꿔서 작성하기만 하면된다.

관건은 위치를 바꾸는 동작을 할 때 count를 증가 시키다가 count가 K와 같아질 때 return을 해야 한다는 것이다.

 

더이상 바꿀 노드가 없을 때 까지만 재귀함수가 동작하기 때문에 바꿀 노드가 없다면 함수를 알아서 빠져 나올 것이다.

그럼에도 count값이 K에 도달하지 못한 것이라면 두 번째 예제와 같은 상황이기 때문에 -1을 출력해야 한다.

 

✨ 풀이

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

public class Main {
    public static int count = 0;
    public static int k = 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()); // n 입력
        k = Integer.parseInt(st.nextToken()); // k 입력

        int[] A = new int[n + 1]; // root는 0으로 채워서 비워놔야 한다.

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i < A.length; i++) {
            A[i] = Integer.parseInt(st.nextToken()); // A배열 채우기
        }

        heapSort(A, n); // 함수를 호출해도 더 이상 바꿀 게 없다면 그냥 빠져나옴

        if (!flag) {
            /*
             * heapSort 함수를 빠져나왔다는 것은
             * heap의 정렬이 모두 끝난것을 의미
             * 정렬이 완료 되었어도 k보다 작다면
             * -1을 출력
             */
            System.out.println(-1);
        }

    }

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

    public static void swap(int[] A, int a, int b) {
        count++; // swap 할때 count 증가

        int temp = A[a];
        A[a] = A[b];
        A[b] = temp;

        if (count == k) { // count하다가 k수 만큼 교환을 했다면
            StringBuilder sb = new StringBuilder();
            for (int i = 1; i < A.length; i++) {
                sb.append(A[i]);
                sb.append(" ");
            }
            System.out.println(sb);
            flag = true; // 원하는 횟수만큼 교환했거나 정렬을 마쳐서 더이상 교환할 필요가 없기 때문에 true로 바꿔줌
        }
    }

    public static void buildMinHeap(int[] A, int n) {
        for (int i = n / 2; i > 0; 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 smaller = -1;

        if (right <= n) {
            smaller = A[left] < A[right] ? left : right;
        } else if (left <= n) {
            smaller = left;
        } else {
            return;
        }

        if (A[smaller] < A[k]) {
            swap(A, k, smaller);
            heapify(A, smaller, n);
        }
    }
}

 

🪲 시간초과

처음 정답 제출을 할 때 Scanner와 출력문은 String에 '+'연산으로 문자열을 이어붙여 출력문을 만들었다.

시간 초과가 뜬 것을 보고 백준 문제풀이를 할 때 Scanner사용에 주의해야 한다고 얼핏 본 기억이 있다.

 

입력은 공백과 '\n'으로 구분하여 입력해야 하기 때문에

Scanner 대신 BufferedReader 를 쓰고 StringTokenizer을 이용해 숫자 하나씩 받는 코드로 바꿨다.

 

출력문은 문자열을 '+'연산으로 이어붙여 출력하는 것이 아닌

StringBuilder를 사용해 append하였다. (StringBuffer 도 가능)

 

원래코드

package baekjoon.n24174;

import java.util.Scanner;

public class Main {
    public static int count = 0;
    public static int k = 0;
    public static boolean flag = false;

    public static void main(String[] args) throws IOException {
    	Scanner scanner = new Scanner(System.in);

        int n = scanner.nextInt(); // n 입력
        k = scanner.nextInt(); // k 입력

        int[] A = new int[n + 1]; // root는 0으로 채워서 비워놔야 한다.

        for (int i = 1; i < A.length; i++) {
            A[i] = scanner.nextInt(); // A배열 채우기
        }

        heapSort(A, n); // 함수를 호출해도 더 이상 바꿀 게 없다면 그냥 빠져나옴

        if (!flag) {
            /*
             * heapSort 함수를 빠져나왔다는 것은
             * heap의 정렬이 모두 끝난것을 의미
             * 정렬이 완료 되었어도 k보다 작다면
             * -1을 출력
             */
            System.out.println(-1);
        }

    }

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

    public static void swap(int[] A, int a, int b) {
        count++; // swap 할때 count 증가

        int temp = A[a];
        A[a] = A[b];
        A[b] = temp;

        if (count == k) { // count하다가 k수 만큼 교환을 했다면
            for (int i = 1; i < A.length; i++) {
            	System.out.print(A[i] + " ");
            }
            flag = true; // 원하는 횟수만큼 교환했거나 정렬을 마쳐서 더이상 교환할 필요가 없기 때문에 true로 바꿔줌
        }
    }

    public static void buildMinHeap(int[] A, int n) {
        for (int i = n / 2; i > 0; 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 smaller = -1;

        if (right <= n) {
            smaller = A[left] < A[right] ? left : right;
        } else if (left <= n) {
            smaller = left;
        } else {
            return;
        }

        if (A[smaller] < A[k]) {
            swap(A, k, smaller);
            heapify(A, smaller, n);
        }
    }
}