;; Indexable Heap
;; MIT - AlexJGriffith - 2025

;; Example:
;; (local heap (require :heap))
;; (local h (heap.new))
;; (local obj1 {:name :obj1})
;; (local obj2 {:name :obj2})
;; (local obj3 {:name :obj3})
;; (h:insert obj1 3)
;; (h:insert obj2 2)
;; (h:insert obj3 5)
;; (h:delete-obj obj2)
;; (h:update-obj-value obj1 7)
;; (h:get-min)
;; (h:pop-min)
;; (h:pop-min)

(local heap-mt {})

(fn parent [index]
  (+ (math.floor (/ (- index 1) 2)) 1))

(fn left [index]
  (+ (* (- index 1) 2) 2))

(fn right [index]
  (+ (* (- index 1) 2) 3))

(fn heapify-down [tab index?]
  (when (> (# tab) 0)
    (var index (or index? 1))
    (var value (?. tab index :value))
    (var obj (?. tab index :obj))
    (var il (left index))
    (var l (or (?. tab il :value) value))
    (var ir (right index))
    (var r (or (?. tab ir :value) value))
    (while (or (> value l) (> value r))
      (if (< l r)
          (let [t (. tab index)
                other (. tab il)]
            (tset tab obj il)
            (tset tab other.obj index)
            (tset tab index other)
            (tset tab il t)
            (set index il))
          (let [t (. tab index)
                other (. tab ir)]
            (tset tab obj ir)
            (tset tab other.obj index)
            (tset tab index other)
            (tset tab ir t)
            (set index ir)))
      (set value (. tab index :value))
      (set obj (. tab index :obj))
      (set il (left index))
      (set l (or (?. tab il :value) value))
      (set ir (right index))
      (set r (or (?. tab ir :value) value)))))

(fn heapify-up [tab index?]
  (when (> (# tab) 0)
    (var index (or index? (# tab)))
    (var value (?. tab index :value))
    (var obj (?. tab index :obj))
    (var ip (parent index))
    (var p (or (?. tab ip :value) value))
    (while (< value p) 
      (let [t (. tab index)
            other (. tab ip)]
        (tset tab obj ip)
        (tset tab other.obj index)
        (tset tab index other)
        (tset tab ip t)
        (set index ip))
      (set value (?. tab index :value))
      (set obj (?. tab index :obj))
      (set ip (parent index))
      (set p (or (?. tab ip :value) value)))))

;; API
(fn heap-mt.update-obj-value [tab obj value]
  (assert (. tab obj) (string.format "Object not in table."))
  (let [index (. tab obj)
        old-value (. tab index :value)]
    (tset tab index :value value)
    (if (< value old-value)
        (heapify-up tab index)
        (heapify-down tab index))))

(fn heap-mt.insert [tab obj value]
  (assert (not (. tab obj)) (string.format "Object already in table."))
  (table.insert tab {: value : obj})
  (tset tab obj (# tab))
  (heapify-up tab))

(fn heap-mt.delete-obj [tab obj]
  (assert (. tab obj) (string.format "Object not in table."))
  (local index (. tab obj))
  (tset tab (. tab (# tab) :obj) index)
  (tset tab obj nil)
  (tset tab index (. tab (# tab)))
  (tset tab (# tab) nil)
  (when (. tab index)
    (heapify-down tab index))
  tab)

(fn heap-mt.get-min [tab]
  (. tab 1 :obj))

(fn heap-mt.pop-min [tab]
  (when (> (# tab) 0)
    (let [ret (. tab 1 :obj)]
      (heap-mt.delete-obj tab ret)
      ret)))

(fn new-heap []
  (setmetatable [] {:__index heap-mt}))

{:new new-heap}

Generated by alexjgriffith using scpaste at Wed Apr 30 23:10:58 2025. EDT. (original)