11286번: 절댓값 힙

풀이

절댓값 힙의 경우 일반적인 최소힙을 사용하는 경우 절댓값으로 만들기 전 값을 알 수 없으므로 절댓값과 원래 값을 묶어서 저장해야 합니다.

Pair 구조체 만들기

절댓값과 원본값을 가지고 있는 구조체를 만듭니다. 해당 구조체는 Comparable Protocol을 준수해야 하는데요. 이 프로토콜은 func <를 구현함으로서 준수할 수 있습니다.

struct Pair: Comparable {
    let absolute: Int
    let origin: Int

    init(_ n: Int) {
        self.absolute = abs(n)
        self.origin = n
    }

    static func < (lhs: Pair, rhs: Pair) -> Bool { //👉 lhs가 rhs보다 작을 때 true
        if lhs.absolute == rhs.absolute { // 절댓값이 같을 때는
            return lhs.origin < rhs.origin // 원본값을 비교
        } else {
            return lhs.absolute < rhs.absolute // 절댓값 비교
        }
    }
}

절댓값 힙 구현

element가 Pair라는 점을 제외하면 Swift로 구현한 다른 최소힙과 완전히 동일합니다. Pair가 Comparable Protocol을 준수하기 때문에 그냥 일반 부등호로 비교해도 문제 없습니다. 최소힙에 대한 자세한 설명이 필요하신 분은 이 포스팅을 참고해주세요.

struct AbsoluteHeap {
    var heap = [Pair]()
    
    mutating func push(_ pair: Pair) {
        if heap.isEmpty {
            heap.append(pair)
            heap.append(pair)
            return
        }
        
        func isMoveUp(_ insertIndex: Int) -> Bool {
            if insertIndex <= 1 {
                return false
            }
            let parentIndex = insertIndex / 2
            return heap[insertIndex] < heap[parentIndex]
        }
        
        heap.append(pair)
        var insertIndex = heap.count - 1
        
        while isMoveUp(insertIndex) {
            let parentIndex = insertIndex / 2
            heap.swapAt(insertIndex, parentIndex)
            insertIndex = parentIndex
        }
    }
    
    enum moveDownDirection { case none, left, right }
    
    mutating func pop() -> Pair? {
        guard heap.count > 1 else { return nil }
        
        func moveDown(_ poppedIndex: Int) -> moveDownDirection {
            let leftChildIndex = poppedIndex * 2
            let rightChildIndex = leftChildIndex + 1
            
            if leftChildIndex >= heap.count {
                return .none
            }
            
            if rightChildIndex >= heap.count {
                return heap[poppedIndex] < heap[leftChildIndex] ? .none : .left
            }
            
            if (heap[poppedIndex] < heap[leftChildIndex]) && (heap[poppedIndex] < heap[rightChildIndex]) {
                return .none
            }
            
            if (heap[poppedIndex] > heap[leftChildIndex]) && (heap[poppedIndex] > heap[rightChildIndex]) {
                return heap[leftChildIndex] < heap[rightChildIndex] ? .left : .right
            }
            
            if heap[poppedIndex] > heap[leftChildIndex] || heap[poppedIndex] > heap[rightChildIndex] {
                return heap[leftChildIndex] < heap[rightChildIndex] ? .left : .right
            }
            
            return .none
        }
        
        let returnData = heap[1]
        heap.swapAt(1, heap.count - 1)
        heap.removeLast()
        var poppedIndex = 1
        
        while true {
            switch moveDown(poppedIndex) {
            case .none:
                return returnData
            case .left:
                let leftChildIndex = poppedIndex * 2
                heap.swapAt(poppedIndex, leftChildIndex)
                poppedIndex = leftChildIndex
            case .right:
                let rightChildIndex = poppedIndex * 2 + 1
                heap.swapAt(poppedIndex, rightChildIndex)
                poppedIndex = rightChildIndex
            }
        }
    }
}

정답 출력

N만큼 반복문을 돌면서 0일 경우에 pop해서 출력하고 그렇게 않을 경우에는 push합니다

let N = Int(readLine()!)!
var absoluteHeap = AbsoluteHeap()

for _ in 0..<N {
    let input = Int(readLine()!)!
    if input == 0 {
        let pair = absoluteHeap.pop()
        pair != nil ? print(pair!.origin) : print(0) //👉 힙이 비어서 nil이 pop되면 0 출력
    } else {
        absoluteHeap.push(Pair(input))
    }
}

😅  그냥 튜플을 이용해도 됩니다.

튜플의 element들이 모두 Comparable을 준수한다면 튜플 자체도 Comparable을 준수하는 자료형입니다. 예를 들어 이번 풀이에 사용되는 (Int, Int) 튜플의 경우 일단 첫 번째 element를 비교하고 만약에 같다면 두 번째 element를 비교합니다. 따라서 우리가 만든 구조체인 Pair와 정확히 일치하게 대소비교를 합니다.

typealias Pair = (Int, Int)

struct AbsoluteHeap {
    var heap = [Pair]()
    
    mutating func push(_ pair: Pair) {
        if heap.isEmpty {
            heap.append(pair)
            heap.append(pair)
            return
        }
        
        func isMoveUp(_ insertIndex: Int) -> Bool {
            if insertIndex <= 1 {
                return false
            }
            let parentIndex = insertIndex / 2
            return heap[insertIndex] < heap[parentIndex]
        }
        
        heap.append(pair)
        var insertIndex = heap.count - 1
        
        while isMoveUp(insertIndex) {
            let parentIndex = insertIndex / 2
            heap.swapAt(insertIndex, parentIndex)
            insertIndex = parentIndex
        }
    }
    
    enum moveDownDirection { case none, left, right }
    
    mutating func pop() -> Pair? {
        guard heap.count > 1 else { return nil }
        
        func moveDown(_ poppedIndex: Int) -> moveDownDirection {
            let leftChildIndex = poppedIndex * 2
            let rightChildIndex = leftChildIndex + 1
            
            if leftChildIndex >= heap.count {
                return .none
            }
            
            if rightChildIndex >= heap.count {
                return heap[poppedIndex] < heap[leftChildIndex] ? .none : .left
            }
            
            if (heap[poppedIndex] < heap[leftChildIndex]) && (heap[poppedIndex] < heap[rightChildIndex]) {
                return .none
            }
            
            if (heap[poppedIndex] > heap[leftChildIndex]) && (heap[poppedIndex] > heap[rightChildIndex]) {
                return heap[leftChildIndex] < heap[rightChildIndex] ? .left : .right
            }
            
            if heap[poppedIndex] > heap[leftChildIndex] || heap[poppedIndex] > heap[rightChildIndex] {
                return heap[leftChildIndex] < heap[rightChildIndex] ? .left : .right
            }
            
            return .none
        }
        
        let returnData = heap[1]
        heap.swapAt(1, heap.count - 1)
        heap.removeLast()
        var poppedIndex = 1
        
        while true {
            switch moveDown(poppedIndex) {
            case .none:
                return returnData
            case .left:
                let leftChildIndex = poppedIndex * 2
                heap.swapAt(poppedIndex, leftChildIndex)
                poppedIndex = leftChildIndex
            case .right:
                let rightChildIndex = poppedIndex * 2 + 1
                heap.swapAt(poppedIndex, rightChildIndex)
                poppedIndex = rightChildIndex
            }
        }
    }
}

let N = Int(readLine()!)!
var absoluteHeap = AbsoluteHeap()

for _ in 0..<N {
    let input = Int(readLine()!)!
    if input == 0 {
        let pair = absoluteHeap.pop()
        pair != nil ? print(pair!.1) : print(0)
    } else {
        absoluteHeap.push((abs(input), input))
    }
}