AoC Smalltalk Heap.st module

" Heap class to implement priority queue "
Object subclass: Heap [
    | arr end |
    Heap class >> new       [ ^super new init: 20000 ]
    Heap class >> new: size [ ^super new init: size  ]

    init: size [
        arr := Array new: size.
        end := 0.
        ^self
    ]

    isEmpty  [ ^(end  = 0) ]
    notEmpty [ ^(end ~= 0) ]

    heapify: idx [
        | left right smallest |
        left  := 2 * idx.
        right := 2 * idx + 1.

        smallest := idx.

        ((left <= end) and: [(arr at: idx) > (arr at: left)]) ifTrue: [
            smallest := left
        ].

        ((right <= end) and: [(arr at: smallest) > (arr at: right)]) ifTrue: [
            smallest := right
        ].

        (smallest ~= idx) ifTrue: [
            arr swap: idx with: smallest.
            self heapify: smallest.
        ]
    ]

    insert: item [
        | i parent |
        end := end + 1.
        (end > arr size) ifTrue: [
            stderr nextPutAll: ('Heap size (%1) exceeded!' % {arr size}); nl.
            Error signal.
        ].

        arr at: end put: item.

        i := end.
        [(i ~= 1) and: [(arr at: (parent := i // 2)) > (arr at: i)]] whileTrue: [
            arr swap: i with: parent.
            i := parent.
        ].
        ^item
    ]

    next [
        | top |
        (end = 0) ifTrue: [ ^nil ].

        top := arr first.
        end := end - 1.

        (end > 0) ifTrue: [
            arr at: 1 put: (arr at: end + 1).
            self heapify: 1.
        ].
        ^top
    ]

    peek [
        ^(end = 0) ifTrue: [nil] ifFalse: [arr first]
    ]

    printOn: aStream [
        (end = 0) ifTrue: [
            aStream nextPutAll: 'Heap (empty)'.
        ] ifFalse: [
            aStream nextPutAll: 'Heap (size: %1, top: %2)' % {end. arr first}.
        ]
    ]
]