프로그래머스 - 디스크 컨트롤러 (Lv.3)

문제 링크

풀이

import Foundation

func solution(_ jobs:[[Int]]) -> Int {
    var jopArr = jobs.map { Jop(start: $0[0], duration: $0[1]) }
    jopArr = jopArr.sorted(by: <)   // 작업 시작시간 순으로 정렬

    var finish = 0
    var result = 0
    var heap = Heap()
    
    while let jop = jopArr.first {
        if finish <= jop.start {
            if let deqeueu = heap.dequeue() {
                result += finish - deqeueu.start + deqeueu.duration
                finish += deqeueu.duration
            } else {
                jopArr.removeFirst()
                result += jop.duration
                finish += jop.start - finish + jop.duration
            }
        } else {
            jopArr.removeFirst()
            heap.enqueue(jop)
        }  
    }

    while let deqeue = heap.dequeue() {
        result += finish - deqeue.start + deqeue.duration
        finish += deqeue.duration
    }
    
    return result / jobs.count
}

struct Jop: Comparable {
    let start: Int
    let duration: Int
    
    static func < (lhs: Jop, rhs: Jop) -> Bool {
        if lhs.start == rhs.start { // 시작시간이 같다면 소요시간으로 구별
            return lhs.duration < rhs.duration
        } else {
            return lhs.start < rhs.start
        }
    }
}

// Jop의 작업시간이 적은 것부터 정렬된 힙
struct Heap {
    var array: [Jop] = []
    
    mutating func enqueue(_ jop: Jop) {
        array.append(jop)
        shiftUp(array.count - 1)
    }
    
    mutating func dequeue() -> Jop? {
        guard !array.isEmpty else { return nil }
        array.swapAt(0, array.count - 1)
        
        defer { shiftDown(0) }
        return array.removeLast()
    }
    
    private func parentIndex(_ childIndex: Int) -> Int {
        return (childIndex - 1) / 2
    }
    
    private func leftChildIndex(_ parentIndex: Int) -> Int {
        return parentIndex * 2 + 1
    }
    
    private func rightChildIndex(_ parentIndex: Int) -> Int {
        return parentIndex * 2 + 2
    }
    
    mutating func shiftUp(_ childIndex: Int) {
        var childIdx = childIndex
        var parentIdx = parentIndex(childIndex)
        
        while childIdx > 0, array[childIdx].duration < array[parentIdx].duration {
            array.swapAt(childIdx, parentIdx)
            childIdx = parentIdx
            parentIdx = parentIndex(childIdx)
        }
    }
    
    mutating func shiftDown(_ parentIndex: Int) {
        var parentIdx = parentIndex
        while true {
            let leftIdx = leftChildIndex(parentIdx)
            let rightIdx = rightChildIndex(parentIdx)
            var candidate = parentIdx
            
            if array.count > rightIdx && array[candidate].duration > array[rightIdx].duration {
                candidate = rightIdx
            } 
            
            if array.count > leftIdx && array[candidate].duration > array[leftIdx].duration {
                candidate = leftIdx
            } 
            
            if candidate == parentIdx {
                break
            }
            
            array.swapAt(parentIdx, candidate)
            parentIdx = candidate
        }
    }
}

후기

힙을 이용해서 풀어야하는 문제.
Swift에 힙이 없기에 만약 다른 언어랑 같이 경쟁 해야하는 테스트라면 많이 억울할 거 같다. (예를 들어 파이썬...)
공부했었던 힙을 구상하며 다른 힙 예제들 안보고 직접 만들고 푸는터라 자꾸 오류가 생겨서 고생했었다.
힙만 잘 쓴다면 어려울 것 없는 문제.