Swift) 백준 14235번 크리스마스 선물 - S3

https://www.acmicpc.net/problem/14235

14235번: 크리스마스 선물

항상 swift로 코테 풀때마다 느끼지만.. 자료구조좀 지원해 줬으면 한다... 솔직히 코테때 이걸 구현하는데 시간 쓸듯...

우선순위 큐만 구현하면 쉽게 풀수 있는 문제여서 따로 설명은 안하겠다.

import Foundation

let n = Int(readLine()!)!
struct Heap<T: Comparable> {
    private var elements: [T]
    private var sortFunction: (T, T) -> Bool
    
    var isEmpty: Bool {
        return elements.isEmpty
    }
    
    var peek: T? {
        if self.isEmpty {
            return nil
        }
        return self.elements[0]
    }
    
    var count: Int {
        return elements.count
    }
    
    init(elements: [T], sortFunction: @escaping (T, T) -> Bool) {
        self.elements = elements
        self.sortFunction = sortFunction
        buildHeap()
    }
    
    private mutating func buildHeap() {
        for i in (0..<count/2).reversed() {
            siftDown(fromIndex: i)
        }
    }
    
    private mutating func siftDown(fromIndex index: Int) {
        var parentIndex = index
        while true {
            let leftChildIndex = 2 * parentIndex + 1
            let rightChildIndex = 2 * parentIndex + 2
            var highestPriorityIndex = parentIndex
            
            if leftChildIndex < count && sortFunction(elements[leftChildIndex], elements[highestPriorityIndex]) {
                highestPriorityIndex = leftChildIndex
            }
            
            if rightChildIndex < count && sortFunction(elements[rightChildIndex], elements[highestPriorityIndex]) {
                highestPriorityIndex = rightChildIndex
            }
            
            if highestPriorityIndex == parentIndex {
                break
            }
            elements.swapAt(parentIndex, highestPriorityIndex)
            parentIndex = highestPriorityIndex
        }
    }
    
    private mutating func siftUp(fromIndex index: Int) {
        var childIndex = index
        var parentIndex = (childIndex - 1) / 2
        while childIndex > 0 && sortFunction(elements[childIndex], elements[parentIndex]) {
            elements.swapAt(childIndex, parentIndex)
            childIndex = parentIndex
            parentIndex = (childIndex - 1) / 2
        }
    }
    
    mutating func insert(_ element: T) {
        elements.append(element)
        siftUp(fromIndex: count - 1)
    }
    
    mutating func pop() -> T? {
        guard !isEmpty else { return nil }
        elements.swapAt(0, count-1)
        let root = elements.removeLast()
        siftDown(fromIndex: 0)
        return root
    }
}

var heap = Heap<Int>(elements: []) { $0 > $1 }

for _ in 0..<n {
    var aArr = readLine()!.split(separator: " ").map { Int($0)! }
    let num = aArr.removeFirst()
    if num == 0 {
        print("\(heap.pop() ?? -1)")
    } else {
        for i in 0..<num {
            heap.insert(aArr[i])
        }
    }
    
}