틈틈히 적어보는 개발 일기

[TIL #5] MaxHeap, 최대힙 본문

📝 TIL

[TIL #5] MaxHeap, 최대힙

itllbegone 2022. 4. 18. 18:08
import Foundation

// Node = (index, value)
typealias Node = (Int, Int)
struct MaxHeap {
    private var nodes: [Int] = [-1]
    
    func print() {
        var from = 1
        var to = 2
        while to < nodes.count {
            Swift.print(nodes[from..<to])
            from = to
            to *= 2
        }
        Swift.print(nodes[from...], terminator: "\n\n")
    }
    
    mutating func insert(value: Int) -> Node {
        nodes.append(value)
        let node = (nodes.count-1, value)
        arrangeUp(newNode: node)
        
        return node
    }
    
    mutating func remove() -> Int? {
        if nodes.count == 1 { return nil }
        nodes.swapAt(1, nodes.count-1)
        let res = nodes.removeLast()
        arrangeDown()
        
        return res
    }

    private func parent(from index: Int) -> Node? {
        if index/2 == 0 { return nil }
        return (index/2, nodes[index/2])
    }
    
    private func leftChild(from index: Int) -> Node? {
        if index <= 0 || index*2 >= nodes.count { return nil }
        return (index*2, nodes[index*2])
    }
    
    private func rightChild(from index: Int) -> Node? {
        if index <= 0 || index*2+1 >= nodes.count { return nil }
        return (index*2+1, nodes[index*2+1])
    }
    
    private mutating func arrangeUp(newNode: Node) {
        var newNode = newNode
        while let parent = parent(from: newNode.0) {
            if parent.1 < newNode.1 {
                nodes.swapAt(parent.0, newNode.0)
                newNode.0 = parent.0
            } else {
                break
            }
        }
    }
    
    private mutating func arrangeDown() {
        if nodes.count == 1 { return }
        var node = (1, nodes[1])
        while let left = leftChild(from: node.0) {
            let right = rightChild(from: node.0) ?? (-1, Int.min)
            let bigger = left.1 < right.1 ? right : left
            if node.1 < bigger.1 {
                nodes.swapAt(node.0, bigger.0)
                node.0 = bigger.0
            } else {
                break
            }
        }
    }
}

var maxHeap = MaxHeap()
maxHeap.insert(value: 1)
maxHeap.insert(value: 5)
maxHeap.insert(value: 3)
maxHeap.insert(value: 4)
maxHeap.insert(value: 2)
maxHeap.insert(value: 1)
maxHeap.insert(value: 2)
maxHeap.print()
maxHeap.insert(value: 7)
maxHeap.print()
maxHeap.remove()
maxHeap.print()
maxHeap.remove()
maxHeap.print()
/*
 result:
 [5]
 [4, 3]
 [1, 2, 1, 2]

 [7]
 [5, 3]
 [4, 2, 1, 2]
 [1]

 [5]
 [4, 3]
 [1, 2, 1, 2]

 [4]
 [2, 3]
 [1, 2, 1]
 */
Comments