백준 1927, 11279 최소 최대 힙

출처: https://www.acmicpc.net/problem/1927 , https://www.acmicpc.net/problem/11279

분류: 자료구조, 힙

---

접근방식

최대 or 최소를 최상위에 두는 heap 을 사용하는 문제였습니다.

복습하는 느낌으로 필요한 부분만 직접 구현해봤습니다 :)

 

해결방법

heap 을 생성할 때 기준을 정해줘서 최대, 최소 힙을 만들 수 있습니다.

heap<Int>(by: <) // 최소힙

import Foundation

struct Heap<T: Comparable> {
    var nodes = [T]()
    var criteria: (T, T) -> Bool
    
    init(by criteria: @escaping (T, T) -> Bool) {
        self.criteria = criteria
    }
    
    mutating func insert(_ value: T) {
        nodes.append(value)
        shiftUp(from: nodes.count-1)
    }
    
    mutating func remove() -> T? {
        guard !nodes.isEmpty else { return nil }
        nodes.swapAt(0, nodes.count-1)
        let min = nodes.removeLast()
        shiftDown(from: 0)
        return min
    }
    
    mutating func shiftDown(from start: Int) {
        let left = start * 2 + 1
        let right = start * 2 + 2
        
        var parent = start
        if left < nodes.count, criteria(nodes[left], nodes[parent]) {
            parent = left
        }
        if right < nodes.count, criteria(nodes[right], nodes[parent]) {
            parent = right
        }
        
        guard parent != start else { return }
        nodes.swapAt(parent, start)
        shiftDown(from: parent)
    }
    
    mutating func shiftUp(from start: Int) {
        var child = start
        var parent = (start-1) / 2
        
        while parent >= 0, criteria(nodes[child], nodes[parent]) {
            nodes.swapAt(child, parent)
            child = parent
            parent = (child-1) / 2
        }
    }
}

let n = Int(readLine()!)!
var heap = Heap<Int>(by: >)

for _ in 0..<n {
    let x = Int(readLine()!)!
    guard x != 0 else {
        print(heap.remove() ?? 0)
        continue
    }

    heap.insert(x)
}