백준 - 수 묶기 (1744)

import Foundation

public struct Heap<T> {
    
    private var nodes = [T]()
    
    private var orderCriteria: (T, T) -> Bool
    
    public init(sort: @escaping (T, T) -> Bool) {
        self.orderCriteria = sort
    }
    
    public init(array: [T], sort: @escaping (T, T) -> Bool) {
        self.orderCriteria = sort
        configureHeap(from: array)
    }
    
    public var count: Int {
        return nodes.count
    }
    
    public func peek() -> T? {
        return nodes.first
    }
    
    func isEmpty() -> Bool {
        return nodes.isEmpty
    }
    
    public mutating func insert(_ value: T) {
        nodes.append(value)
        shiftUp(nodes.count - 1)
    }
    
    public mutating func remove() -> T? {
        guard !nodes.isEmpty else { return nil }
        
        if nodes.count == 1 {
            return nodes.removeLast()
        } else {
            let value = nodes[0]
            nodes[0] = nodes.removeLast()
            shiftDown(0)
            return value
        }
    }
    
    public mutating func remove(at index: Int) -> T? {
        guard index < nodes.count else { return nil }
        
        let lastIndex = nodes.count-1
        if index != lastIndex {
            nodes.swapAt(index, lastIndex)
            shiftDown(from: index, until: lastIndex)
            shiftUp(index)
        }
        
        return nodes.removeLast()
    }
    
    private mutating func configureHeap(from array: [T]) {
           nodes = array
           
           for i in stride(from: nodes.count/2 - 1, through: 0, by: -1) {
               shiftDown(i)
           }
       }
    
    private func parentIndex(ofIndex i: Int) -> Int {
        return (i - 1) / 2
    }
    
    private func leftChildIndex(ofIndex i: Int) -> Int {
        return 2*i + 1
    }
    
    private func rightChildIndex(ofIndex i: Int) -> Int {
        return 2*i + 2
    }
    
    private mutating func shiftUp(_ index: Int) {
        var childIndex = index
        let child = nodes[childIndex] // 처음에 노드를 저장해두고 인덱스를 구한 후 바꿔준다.
        var parentIndex = self.parentIndex(ofIndex: index)
        
        while childIndex > 0 && orderCriteria(child, nodes[parentIndex]) {
            nodes[childIndex] = nodes[parentIndex]
            childIndex = parentIndex
            parentIndex = self.parentIndex(ofIndex: childIndex)
        }
        
        nodes[childIndex] = child
    }
    
    private mutating func shiftDown(from index: Int, until endIndex: Int) {
        let leftChildIndex = self.leftChildIndex(ofIndex: index)
        let rightChildIndex = leftChildIndex + 1
        
        var first = index
        if leftChildIndex < endIndex && orderCriteria(nodes[leftChildIndex], nodes[first]) {
            first = leftChildIndex
        }
        if rightChildIndex < endIndex && orderCriteria(nodes[rightChildIndex], nodes[first]) {
            first = rightChildIndex
        }
        if first == index { return }
        
        nodes.swapAt(index, first)
        shiftDown(from: first, until: endIndex)
    }
    
    private mutating func shiftDown(_ index: Int) {
        shiftDown(from: index, until: nodes.count)
    }
    
}

var heap = Heap<(Int, Int)>(array: [], sort: { $0.0 < $1.0 })

let N = Int(readLine()!)!
var plusQueue = Heap<Int>(sort: >)
var minusQueue = Heap<Int>(sort: <)
var zero = 0
var one = 0
var answer = 0

for _ in 0..<N {
    let num = Int(readLine()!)!
    if num > 1 {
        plusQueue.insert(num)
    } else if num == 1 {
        one += 1
    } else if num == 0 {
        zero += 1
    } else {
        minusQueue.insert(num)
    }
}

while plusQueue.count >= 2 {
    let num1 = plusQueue.remove()!
    let num2 = plusQueue.remove()!
    answer += num1 * num2
}

if !plusQueue.isEmpty() {
  answer += plusQueue.remove()!
}

while minusQueue.count >= 2 {
    let num1 = minusQueue.remove()!
    let num2 = minusQueue.remove()!
    answer += num1 * num2
}

if !minusQueue.isEmpty() {
    let num = minusQueue.remove()!
    if zero > 0 {
        zero -= 1
    } else {
        answer += num
    }
}

answer += one

print(answer)

• 수의 집합을 1보다 큰 수 , 1 , 0 , 음수 4가지 그룹으로 나눈다.
• 1보다 큰 수 그룹은 우선순위 큐를 사용해 내림차순으로 정렬한다.
• 음수 그룹은 우선순위 큐를 사용해 오름차순으로 정렬한다.

→ 이렇게 정렬해야 가능한 큰 수들끼리 곱해서 결과값을 크게 만들 수 있다.

• 1보다 큰 수 그룹에 있는 수끼리 차례대로 곱한 후 더한다.
  • 남은 값이 있다면 그대로 결과값에 더해준다.
• 음수 그룹에 있는 수끼리 차례대로 곱한 후 더한다.
  • 남은 값이 있는 경우, 0이 수열에 있다면 0과 곱해서 0을 만들고 그렇지 않다면 그 값 그대로 결과값에 더해준다.
• 남은 1들을 모두 결과값에 더해준다.