(24174) 알고리즘 수업 - 힙 정렬 2

문제

오늘도 서준이는 최소 힙 기반 힙 정렬 수업 조교를 하고 있다. 아빠가 수업한 내용을 학생들이 잘 이해했는지 문제를 통해서 확인해보자.
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 ≤ 108)가 주어진다.다음 줄에 서로 다른 배열 A의 원소 A1, A2, ..., AN이 주어진다. (1 ≤ Ai ≤ 109)

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

---

문제 풀이

힙이라는 자료구조를 써서 푸는 첫번째 문제인데 장정 8시간동안 문제를 풀었다... 덕분에 힙 내부에서 어떠한 방식으로 돌아가는지 안보고 쓸 수 있는 정도까지 올라 왔는데, 금방 잊을 듯 하다 ㅎ.. 이렇게 어렵게 문제를 풀었으니 더욱 상세하게 정리 해보고자 한다. 

 

 

현재 문제에서 질문하고 있는 것은 최소힙 상태에서 힙정렬을 하는 방식으로 더 이상 정렬이 필요없을 때 까지 반복을 하면 된다. 거기에서 상수 K 값을 주어 졌을 때 해당 K번 만큼 자리 이동이 있었을 때, 배열의 값을 출력해야 하는 문제이다. 

 

해당 문제에서 사용한 함수는 아래와 같다.

private static void heapSort(int[] ints) {
	... // 힙 sort를 수행하는 함수
}

private static void buildHeap(int[] ints, int len) {
    ... // 최소힙을 만드는 함수
}

private static void heapify(int[] ints, int i, int len) {
	... // 자식 노드에서 작은 값을 찾아내는 함수	
}

private static void swap(int[] ints, int i, int targetIdx) {
	... // 부모 노드와 자식 노드의 값을 변경하는 함수
}

 

전부 다 처음 사용하는 함수로써 인터넷에 올라와 있는 글 들을 하나씩 보면서 감을 익혔다. 각각의 메소드에 대해서 설명을 해보자면 아래와 같다. 

 

입력되는 배열을 매개변수로 받아서 길이를 변수 len 으로 선언. 

buildHeap 함수를 호출을 하면서 최초로 최소힙을 생성한다. 

private static void heapSort(int[] ints) {
    int len = ints.length;
    buildHeap(ints, len); // 최소힙을 생성하는 메소드

    for (int i = len - 1; i >= 0; i--) {
        swap(ints, 0, i);
        heapify(ints, 0, i);
    }
}

 

heapSort에서 호출된 bulidHeap 함수이다.

여기서 반복문을 사용하는데 i 의 값을 len /2 -1 로 지정하였다. 이렇게 지정을 한 이유는 자식노드를 가지고 있는 부모노드만 반복문을 수행하도록 한 것으로 하위 트리가 정렬이 되면은 상위 트리로 이동하여 정렬을 한다.

private static void buildHeap(int[] ints, int len) {
    for (int i = len / 2 - 1; i >= 0; i--) {
        heapify(ints, i, len);
    }
}

 

heapify이 함수는 자식 노드중 더 작은 값을 찾는데 사용 된 함수이다. 

자식 노드를 구하는 공식은 아래 코드와 같으며, 배열의 index 0에 노드 값이 아닌 필요 없는 값을 넣게 되었을 때는 아래와 같은 공식을 사용한다. 

배열 index 1 부터 값이 시작할때
leftIdx = cur * 2;
rightIdx = cur * 2 +1;

이번에는 배열inex 0 부터 시작하므로 아래 코드를 이용하였다. 

여기서 더 작은 값을 구하였다면 그것을 바꿔주는 swap 함수를 이용. 그리고 다시 재귀호출로 반복을 하여 조건에 들지 못 할 때 탈출한다. 

private static void heapify(int[] ints, int i, int len) {
    int leftIdx = i * 2 + 1;
    int rightIdx = i * 2 + 2;
    int targetIdx = -1;

    if (rightIdx < len) {
        targetIdx = ints[leftIdx] < ints[rightIdx] ? leftIdx : rightIdx;
    } else if (leftIdx < len) {
        targetIdx = leftIdx;
    } else {
        return;
    }

    if (ints[targetIdx] < ints[i]) {
        swap(ints, i, targetIdx);
        heapify(ints, targetIdx, len);
    }
}

 

swap은 부모노드와 자식노드의 값을 변경해주는 함수이다.

현재 count 값이 LIMIT 값보다 작을 때는 값을 변경해주고, 그렇지 아니 할 때는 배열 값을 출력을 한다.

private static void swap(int[] ints, int i, int targetIdx) {
    if (count++ < LIMIT) {
        int parentVal = ints[i];
        ints[i] = ints[targetIdx];
        ints[targetIdx] = parentVal;


    }
    if (count == LIMIT + 1) {
        StringBuffer sb = new StringBuffer();
        Arrays.stream(ints).forEach(s -> sb.append(s).append(" "));
        System.out.println(sb);
        System.exit(0);
    }

}

---

나의 답안

import java.io.IOException;
import java.util.Arrays;
import java.util.Scanner;

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

    public static void main(String[] args) throws IOException {
        Scanner sc = new Scanner(System.in);
        String[] arr1 = sc.nextLine().split(" ");
        String[] arr2 = sc.nextLine().split(" ");
        int N = Integer.parseInt(arr1[0]);
        LIMIT = Integer.parseInt(arr1[1]);
        int[] ints = Arrays.stream(arr2).mapToInt(Integer::parseInt).toArray();

        heapSort(ints);
        System.out.println(-1);
    }

    private static void heapSort(int[] ints) {
        int len = ints.length;
        buildHeap(ints, len); // 최소힙을 생성하는 메소드

        for (int i = len - 1; i >= 0; i--) {
            swap(ints, 0, i);
            heapify(ints, 0, i);
        }
    }

    private static void buildHeap(int[] ints, int len) {
        for (int i = len / 2 - 1; i >= 0; i--) {
            heapify(ints, i, len);
        }
    }

    private static void heapify(int[] ints, int i, int len) {
        int leftIdx = i * 2 + 1;
        int rightIdx = i * 2 + 2;
        int targetIdx = -1;

        if (rightIdx < len) {
            targetIdx = ints[leftIdx] < ints[rightIdx] ? leftIdx : rightIdx;
        } else if (leftIdx < len) {
            targetIdx = leftIdx;
        } else {
            return;
        }

        if (ints[targetIdx] < ints[i]) {
            swap(ints, i, targetIdx);
            heapify(ints, targetIdx, len);
        }
    }

    private static void swap(int[] ints, int i, int targetIdx) {
        if (count++ < LIMIT) {
            int parentVal = ints[i];
            ints[i] = ints[targetIdx];
            ints[targetIdx] = parentVal;


        }
        if (count == LIMIT + 1) {
            StringBuffer sb = new StringBuffer();
            Arrays.stream(ints).forEach(s -> sb.append(s).append(" "));
            System.out.println(sb);
            System.exit(0);
        }

    }

}