Swift) 프로그래머스 - 이중우선순위큐

https://school.programmers.co.kr/learn/courses/30/lessons/42628?language=swift 

프로그래머스

우선순위큐를 구현 할 수 있다면 쉬운 문제. 오름차순과 내림차순 우선순위 큐를 둘 다 구현하고 명령어에 따라 각 큐를 연산 한 뒤, 연산하지 않은 큐를 새로 만들어 할당하는 식으로 문제를 해결했다.

import Foundation

struct Heap {
    var elements: [Int]
    var sortFunction: (Int, Int) -> Bool
    
    var isEmpty: Bool {
        return elements.isEmpty
    }
    
    var peek: Int? {
        if self.isEmpty {
            return nil
        }
        return self.elements[0]
    }
    
    var count: Int {
        return elements.count
    }
    
    init(_ elements: [Int], _ sortFunction: @escaping(Int, Int) -> Bool) {
        self.elements = elements
        self.sortFunction = sortFunction
        buildHeap()    
    }
    
    mutating func buildHeap() {
        for i in (0..<count/2).reversed() {
            siftDown(fromIndex: i)
        }
    }
    
    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: Int) {
        elements.append(element)
        siftUp(fromIndex: count - 1)
    }
    
    mutating func pop() -> Int? {
        guard !isEmpty else { return nil }
        elements.swapAt(0, count-1)
        let root = elements.removeLast()
        siftDown(fromIndex: 0)
        return root
    }
    
}

func solution(_ operations:[String]) -> [Int] {
    var bigHeap = Heap([]) {$0 < $1}
    var smallHeap = Heap([]) {$0 > $1}
    for operation in operations {
        let oper = operation.split(separator: " ").map { String($0) }
        if oper[0] == "I" {
            bigHeap.insert(Int(oper[1])!)
            smallHeap.insert(Int(oper[1])!)
        } else {
            if Int(oper[1])! == -1 {
                bigHeap.pop()
                smallHeap = Heap(bigHeap.elements) {$0 > $1}
            } else {
                smallHeap.pop()   
                bigHeap = Heap(smallHeap.elements) {$0 < $1}
            }
        }
    }
    if smallHeap.count == 0 {
        return [0,0]
    }
    return [smallHeap.pop()!, bigHeap.pop()!]
}