[C++] 최소 힙 (1927번)

최소 힙

 

문제

널리 잘 알려진 자료구조 중 최소 힙이 있다. 최소 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오.

1. 배열에 자연수 x를 넣는다.
2. 배열에서 가장 작은 값을 출력하고, 그 값을 배열에서 제거한다.

프로그램은 처음에 비어있는 배열에서 시작하게 된다.

입력

첫째 줄에 연산의 개수 N(1 ≤ N ≤ 100,000)이 주어진다. 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 만약 x가 자연수라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0이라면 배열에서 가장 작은 값을 출력하고 그 값을 배열에서 제거하는 경우이다. x는 231보다 작은 자연수 또는 0이고, 음의 정수는 입력으로 주어지지 않는다.

출력

입력에서 0이 주어진 횟수만큼 답을 출력한다. 만약 배열이 비어 있는 경우인데 가장 작은 값을 출력하라고 한 경우에는 0을 출력하면 된다.

 

---

문제풀이

 

priority_queue를 이용하면 순식간에 풀 수 있는 문제다.

하지만 공부할 겸 민 힙을 직접 구현해봤다.

계산상의 편의를 위해 배열은 0번이 아닌 1번 인덱스부터 사용한다.

 

아래의 논리를 코드로 짜보자.

 

<삽입>

원소를 힙의 맨 뒤에 넣는다.

부모 노드와 비교한다.

부모 노드보다 작을 경우 부모 노드와 위치를 바꾼다.

루트 노드에 도달하거나 부모 노드보다 큰 경우 더 이상의 반복을 멈춘다.

 

<삭제>

맨 앞의 원소(루트 노드)를 출력한다.

맨 뒤의 원소를 루트 노드에 넣는다. (이때 힙의 크기를 1 줄여야 한다.)

자식이 존재하지 않을 경우, 가장 작은 자식보다 자신이 작은 경우 -> 끝

자식이 하나 존재할 경우 -> 자식과 자신을 비교하여 작으면 위치를 바꾼다. -> 끝

자식이 둘 존재할 경우 -> 가장 작은 자식과 자신을 비교하여 작으면 위치를 바꾼다. -> 반복

 

#include <iostream>

using namespace std;

int arr[100001];
int idx = 0;

void HeapSort() {
    int s = idx;
    while(s > 1 && arr[s] < arr[s/2]) {
        swap(arr[s], arr[s/2]);
        s /= 2;
    }
}

int HeapPop() {
    int ret = arr[1];

    arr[1] = arr[idx--];

    int now = 1;
    // 자식노드가 존재하는 경우 비교
    while(now*2<=idx) {
        // 자식노드가 하나만 존재하는 경우
        if(now*2+1>idx) {
            if(arr[now]>arr[now*2])
                swap(arr[now], arr[now*2]);
            break;
        }
        // 자식노드가 두개 존재하는 경우
        else {
            // 더 작은 자식노드 선출
            int min_idx;
            if(arr[now*2]<=arr[now*2+1])
                min_idx = now*2;
            else
                min_idx = now*2+1;
            
            // 현재 노드와 더 작은 자식 노드 비교
            if(arr[min_idx] < arr[now]) {
                swap(arr[min_idx], arr[now]);
                now = min_idx;
            }
            else {
                break;
            }
        }
    }
    return ret;
}

int main() {
    ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
    int n;
    cin >> n;
    for(int i=0; i<n; i++) {
        int tmp;
        cin >> tmp;
        if(tmp == 0 && idx == 0) {
            cout << 0 << "\n";
        }
        else if(tmp == 0) {
            cout << HeapPop() << "\n";
        }
        else {
            arr[++idx] = tmp;
            HeapSort();
        }
    }
    return 0;
}