Swift) 백준 17503번 맥주 축제 - S1

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

17503번: 맥주 축제

우선순위 큐를 활용하여 풀면 시간제한 내에 풀 수 있는 문제. 맥주를 도수레벨을 기준으로 오름차순으로 정렬 한 후, 정렬된 맥주를 처음부터 끝까지 탐색한다. 이때 탐색하면서 맥주의 선호도을 기준으로 내림차순으로 정렬하는 우선순위 큐에 넣어주면서, 큐 사이즈가 n보다 클 시 우선순위 큐에서 pop연산을 통해 선호도가 가장 낮은 맥주를 제거하고, 같아지면 반복문을 빠져나오는 식으로 접근했다.

우선순위 큐의 개념을 알고 구현할 수 있다면 간단하게 풀 수 있는 문제.

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
    }
}

struct Beer: Comparable {
    let priority: Int
    let level: Int
    
    static func < (lhs: Beer, rhs: Beer) -> Bool {
        return false
    }
    
}

var beerArr: [Beer] = []

let nmk = readLine()!.split(separator: " ").map{ Int($0)! }
let n = nmk[0]
let m = nmk[1]
let k = nmk[2]

for _ in 0..<k {
    let input = readLine()!.split(separator: " ").map {Int($0)!}
    let beer: Beer = Beer(priority: input[0], level: input[1])
    beerArr.append(beer)
}
beerArr.sort { $0.level < $1.level }
var beerPriority: Heap<Beer> = Heap(elements: []) { $0.priority < $1.priority }

var total: Int = 0
var answer: Int = -1
for beer in beerArr {
    beerPriority.insert(beer)
    total += beer.priority
    if beerPriority.count > n {
        total -= beerPriority.pop()!.priority
    }
    if beerPriority.count == n && total >= m {
        answer = beer.level
        break
    }
}

print(answer)