BOJ 11279 최대 힙

By kuper0201

Published in 알고리즘

2023-05-25

1 min read

문제

입력

풀이

문제

BOJ 11279 최대 힙 바로 가기

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

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

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

입력

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

풀이

해당 문제는 최대 힙을 구현하여 자연수를 추가하고 최대값을 출력하라는 문제이다.

문제를 처음 접하고 가장 먼저 떠올릴 수 있는 방법은 배열에 자연수를 추가 할 때마다 배열을 정렬하는 방식으로 이는 배열의 크기가 커질수록 비효율적인 방법이다.

정렬 알고리즘은 일반적으로 $O(N \log N)$ 의 시간 복잡도를 갖는다.

따라서 배열에 N개의 자연수를 추가할 때마다 정렬을 하게 되면 $O(N^2 \log N)$ 의 시간 복잡도를 가지게 된다.

N의 제한이 100,000이므로 시간 초과를 피할 수 없을 것이다.

하지만 최대 힙을 사용하여 구현할 경우 자연수를 추가하는 작업에 $O(\log N)$ 의 시간 복잡도가 소요되고, 최대값을 구하는 작업에도 $O(\log N)$ 의 시간 복잡도가 소요된다.

따라서 N개의 연산을 모두 수행하는 데에 $O(N \log N)$ 의 시간 복잡도를 가지게 된다.

최대 힙은 다음과 같은 특징을 갖는다.

부모 노드의 값은 항상 자식 노드의 값보다 크거나 같다.
가장 큰 값은 항상 루트 노드에 위치한다.

우선순위 큐는 보통 최소 힙을 이용해 구현되는 자료 구조이다.

최소 힙은 최대 힙과는 반대의 특징을 가지기 때문에 최소 힙에 부호를 바꾼(-1을 곱한) 값을 추가하게 되면 최대 힙과 같은 효과를 가진다.

물론 최대값을 구할 때에도 부호를 바꿔주어야 한다.

대부분의 프로그래밍 언어는 우선순위 큐의 구현을 기본적으로 제공하고 있으며 이를 이용하면 최대 힙을 구현하지 않고 문제를 해결 할 수 있다.

코드 보기(Java)

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

class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();
        
        // 우선순위 큐 생성
        PriorityQueue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());

        // 연산의 개수 입력
        int N = Integer.parseInt(br.readLine());
        for(int i = 0; i < N; i++) {
            // 명령어 입력
            int cmd = Integer.parseInt(br.readLine());

            if(cmd == 0) { // 명령어가 0일 때
                sb.append(pq.isEmpty() ? "0" : pq.poll()); // 비어있으면 0, 비어있지 않으면 최대값 출력
                sb.append("\n");
            } else pq.add(cmd); // 명령어가 0이 아닐 때
        }

        System.out.println(sb);
    }
}

코드 보기(C++)

#include <iostream>
#include <queue>

#define fastio ios_base::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr)
#define endl '\n'

using namespace std;

int main() {
    fastio;
    
    // 우선순위 큐 생성
    priority_queue<int> pq;

    // 연산의 개수 입력
    int N, cmd;
    cin >> N;
    for(int i = 0; i < N; i++) {
        // 명령어 입력
        cin >> cmd;

        if(cmd == 0) { // 명령어가 0일 때
            if(pq.empty()) cout << "0" << endl; // 비어있으면 0
            else { // 비어있지 않으면 최대값 출력
                cout << pq.top() << endl;
                pq.pop();
            }
        } else pq.push(cmd); // 명령어가 0이 아닐 때
    }
    
    return 0;
}

코드 보기(Kotlin)

import java.util.*

fun main(args: Array<String>) = with(System.`in`.bufferedReader()) {
    // 우선순위 큐 생성
    var pq = PriorityQueue<Int>(Collections.reverseOrder())
    
    // 연산의 개수 입력
    var N = readLine().toInt()
    for(i in 0 until N) {
        // 명령어 입력
        var cmd = readLine().toInt()
        if(cmd == 0) { // 명령어가 0일 때
            if(pq.isEmpty()) println("0") // 비어있으면 0
            else println(pq.poll()) // 비어있지 않으면 최대값 출력
        } else pq.add(cmd) // 명령어가 0이 아닐 때
    }
}

코드 보기(Python)

from sys import stdin
import heapq

def main():
    # 우선순위 큐 생성
    pq = []

    # 연산의 개수 입력
    N = int(stdin.readline())
    for i in range(0, N): # 명령어 입력
        cmd = int(stdin.readline())
        if cmd == 0: # 명령어가 0일 때
            if not pq:
                print("0") # 비어있으면 0
            else:
                print(heapq.heappop(pq) * -1) # 비어있지 않으면 최대값 출력
        else: # 명령어가 0이 아닐 때
            heapq.heappush(pq, cmd * -1) # 내림차순 정렬

if __name__ == "__main__":
    main()

코드 보기(Swift)

Swift에서는 우선순위 큐 또는 힙 자료 구조를 기본적으로 지원하지 않기 때문에 직접 구현해야 한다.

우선순위 큐를 구현하여 해당 문제를 해결하는 코드는 아래와 같다.

import Foundation

struct PriorityQueue<T: Comparable> {
    private var elements: [T] = []
    
    var isEmpty: Bool {
        return elements.isEmpty
    }
    
    var count: Int {
        return elements.count
    }
    
    mutating func push(_ element: T) {
        elements.append(element)
        siftUp(from: elements.count - 1)
    }
    
    mutating func poll() -> T? {
        guard !isEmpty else {
            return nil
        }
        
        if elements.count == 1 {
            return elements.removeFirst()
        } else {
            let root = elements[0]
            elements[0] = elements.removeLast()
            siftDown(from: 0)
            return root
        }
    }
    
    private mutating func siftUp(from index: Int) {
        var childIndex = index
        var parentIndex = parent(of: childIndex)
        
        while childIndex > 0 && elements[childIndex] < elements[parentIndex] {
            elements.swapAt(childIndex, parentIndex)
            childIndex = parentIndex
            parentIndex = parent(of: childIndex)
        }
    }
    
    private mutating func siftDown(from index: Int) {
        var parentIndex = index
        
        while true {
            let leftChildIndex = leftChild(of: parentIndex)
            let rightChildIndex = rightChild(of: parentIndex)
            var targetIndex = parentIndex
            
            if leftChildIndex < elements.count && elements[leftChildIndex] < elements[targetIndex] {
                targetIndex = leftChildIndex
            }
            
            if rightChildIndex < elements.count && elements[rightChildIndex] < elements[targetIndex] {
                targetIndex = rightChildIndex
            }
            
            if targetIndex == parentIndex {
                return
            }
            
            elements.swapAt(parentIndex, targetIndex)
            parentIndex = targetIndex
        }
    }
    
    private func parent(of index: Int) -> Int {
        return (index - 1) / 2
    }
    
    private func leftChild(of index: Int) -> Int {
        return 2 * index + 1
    }
    
    private func rightChild(of index: Int) -> Int {
        return 2 * index + 2
    }
}


func main() {
    // 우선순위 큐 생성
    var pq = PriorityQueue<Int>()
    
    // 연산의 개수 입력
    var N = Int(readLine()!)!
    for i in 0..<N {
        // 명령어 입력
        var cmd = Int(readLine()!)!
        if(cmd == 0) { // 명령어가 0일 때
            if pq.isEmpty { // 비어있으면 0
                print("0")
            } else { // 비어있지 않으면 최대값 출력
                print(pq.poll()! * -1)
            }
        } else { // 명령어가 0이 아닐 때
            pq.push(cmd * -1)
        }
    }
}

main()