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

문제: https://programmers.co.kr/learn/courses/30/lessons/42628

코딩테스트 연습 - 이중우선순위큐

 

최소힙과 최대힙 각각 만들어서 계산해주었습니다!

import Foundation

public struct Heap<Element: Equatable> {
    var elements: [Element] = []
    let sort: (Element, Element) -> Bool
    
    init(sort: @escaping (Element, Element) -> Bool, elements: [Element] = []) {
        self.sort = sort
        self.elements = elements
        buildHeap()
    }
    
    var isEmpty: Bool {
        elements.isEmpty
    }
    
    var count: Int {
        elements.count
    }
    
    func peek() -> Element? {
        elements.first
    }
    
    mutating public func buildHeap() {
        if !elements.isEmpty {
            for i in stride(from: elements.count / 2 - 1, through: 0, by: -1) {
                siftDown(from: i)
            }
        }
    }
    
    func leftChildIndex(ofParentAt index: Int) -> Int {
        (2 * index) + 1
    }
    
    func rightChildIndex(ofParentAt index: Int) -> Int {
        (2 * index) + 2
    }
    
    func parentIndex(ofChildAt index: Int) -> Int {
        (index - 1) / 2
    }
    
    mutating func remove() -> Element? {
        guard !isEmpty else {
            return nil
        }
        elements.swapAt(0, count - 1)
        defer {
            siftDown(from: 0)
        }
        return elements.removeLast()
    }
    
    mutating func siftDown(from index: Int) {
        var parent = index
        while true {
            let left = leftChildIndex(ofParentAt: parent)
            let right = rightChildIndex(ofParentAt: parent)
            var candidate = parent
            if left < count && sort(elements[left], elements[candidate]) {
                candidate = left
            }
            if right < count && sort(elements[right], elements[candidate]) {
                candidate = right
            }
            if candidate == parent {
                return
            }
            elements.swapAt(parent, candidate)
            parent = candidate
        }
    }
    
    mutating func insert(_ element: Element) {
        elements.append(element)
        siftUp(from: elements.count - 1)
    }
    
    mutating func siftUp(from index: Int) {
        var child = index
        var parent = parentIndex(ofChildAt: child)
        while child > 0 && sort(elements[child], elements[parent]) {
            elements.swapAt(child, parent)
            child = parent
            parent = parentIndex(ofChildAt: child)
        }
    }
    
    mutating func remove(at index: Int) -> Element? {
        guard index < elements.count else {
            return nil // 1
        }
        if index == elements.count - 1 {
            return elements.removeLast() // 2
        } else {
            elements.swapAt(index, elements.count - 1) // 3
            defer {
                siftDown(from: index) // 5
                siftUp(from: index)
            }
            return elements.removeLast() // 4
        }
    }
    
    func index(of element: Element, startingAt i: Int) -> Int? {
        if i >= count {
            return nil
        }
        if sort(element, elements[i]) {
            return nil
        }
        if element == elements[i] {
            return i
        }
        if let j = index(of: element, startingAt: leftChildIndex(ofParentAt: i)) {
            return j
        }
        if let j = index(of: element, startingAt: rightChildIndex(ofParentAt: i)) {
            return j
        }
        return nil
    }
}

func solution(_ operations:[String]) -> [Int] {
    var minHeap = Heap(sort: { $0 < $1 }, elements: [Int]())
    var maxHeap = Heap(sort: { $0 > $1 }, elements: [Int]())
    
    for op in operations {
        var op = Array(op)
        if op[0] == "I" { // insert
            let str = String(op[2..<op.count])
            let number = Int(str)!
            
            minHeap.insert(number)
            maxHeap.insert(number)
        } else {
            if op[2] == "-" { // D min
                if !minHeap.isEmpty {
                    let min = minHeap.remove()!
                    let index = maxHeap.index(of: min, startingAt: 0)!
                    maxHeap.remove(at: index)
                }
            } else { // D max
                if !maxHeap.isEmpty {
                    let max = maxHeap.remove()!
                    let index = minHeap.index(of: max, startingAt: 0)!
                    minHeap.remove(at: index)
                }
            }
        }
    }
    
    if minHeap.isEmpty {
        return [0, 0]
    } else {
        let min = minHeap.remove()!
        let max = maxHeap.remove()!
        return [max, min]
    }
}

solution(["I 16","D 1"]) // [0, 0]
solution(["I 7","I 5","I -5","D -1"]) // [7,5]