BinaryHeap in Lua

learning-lua/algo/BinaryHeap-test.lua

@@ -0,0 +1,99 @@
+local BinaryHeap = require "algo.BinaryHeap"
+
+local h = BinaryHeap:new()
+h:add('9'):add('8'):add('7'):add('6'):add('5')
+h:verify()
+assert(h:dump() == "5,6,8,9,7")
+assert(h:size() == 5)
+assert(h:remove() == '5')
+h:verify()
+assert(h:dump() == "6,7,8,9")
+assert(h:size() == 4)
+
+assert(h:remove() == '6')
+h:verify()
+assert(h:dump() == "7,9,8")
+assert(h:size() == 3)
+
+assert(h:remove() == '7')
+h:verify()
+assert(h:dump() == "8,9")
+assert(h:size() == 2)
+
+assert(h:remove() == '8')
+h:verify()
+assert(h:dump() == "9")
+assert(h:size() == 1)
+
+assert(h:remove() == '9')
+h:verify()
+assert(h:dump() == "")
+assert(h:size() == 0)
+
+local h = BinaryHeap:new{'9', '8', '7', '6', '5'}
+h:verify()
+assert(h:dump() == "5,6,7,9,8")
+assert(h:size() == 5)
+
+local h = BinaryHeap:new{'9', '8', '7', '6'}
+h:verify()
+assert(h:dump() == "6,8,7,9")
+assert(h:size() == 4)
+
+local h = BinaryHeap:new{3, 1, 2}
+h:verify()
+assert(h:dump() == "1,3,2")
+
+local h = BinaryHeap:new{1, 2, 3, 4, 5}
+h:verify()
+assert(h:dump() == "1,2,3,4,5")
+
+local h = BinaryHeap:new{5, 4, 3, 2, 1}
+h:verify()
+assert(h:dump() == "1,2,3,5,4")
+
+local h = BinaryHeap:new{4, 4, 4, 4}
+h:verify()
+assert(h:dump() == "4,4,4,4")
+
+local h = BinaryHeap:new{3, 1, 9, 8, 4, 6, 7, 5, 2}
+h:verify()
+assert(h:dump() == "1,2,6,3,4,9,7,5,8")
+
+local h = BinaryHeap:new{-1, -3, -2, -4, -5}
+h:verify()
+assert(h:dump() == "-5,-4,-2,-1,-3")
+
+local h = BinaryHeap:newMaxHeap{3, 1, 9, 8, 4, 6, 7, 5, 2}
+h:verify()
+assert(h:dump() == "9,8,7,5,4,6,3,1,2")
+
+local h = BinaryHeap:newMaxHeap()
+h:add('9'):add('8'):add('7'):add('6'):add('5')
+h:verify()
+assert(h:dump() == "9,8,7,6,5", h:dump())
+assert(h:size() == 5)
+assert(h:remove() == '9')
+h:verify()
+assert(h:dump() == "8,6,7,5", h:dump())
+assert(h:size() == 4)
+
+assert(h:remove() == '8')
+h:verify()
+assert(h:dump() == "7,6,5", h:dump())
+assert(h:size() == 3)
+
+assert(h:remove() == '7', h:dump())
+h:verify()
+assert(h:dump() == "6,5", h:dump())
+assert(h:size() == 2)
+
+assert(h:remove() == '6')
+h:verify()
+assert(h:dump() == "5")
+assert(h:size() == 1)
+
+assert(h:remove() == '5')
+h:verify()
+assert(h:dump() == "")
+assert(h:size() == 0)

learning-lua/algo/BinaryHeap.lua

@@ -0,0 +1,127 @@
+-- Reference: https://opendatastructures.org
+-- Courtesy of https://en.wikipedia.org/wiki/J._W._J._Williams
+local BinaryHeap = {}
+
+--- Maintain the heap invariant by repeatedly swapping with the parent if necessary.
+---@param i (number) the starting position; default to the last element.
+function BinaryHeap:bubble_up(i)
+  i = i or #self
+  local a, comp = self, self.comp
+  while i > 1 do
+    local p = i >> 1 -- parent
+    if comp(a[p], a[i]) then -- value stored at index i is not smaller than that of its parent
+      return
+    end
+    a[i], a[p] = a[p], a[i] -- swap value with parent
+    i = p -- repeatedly
+  end
+end
+
+--- Maintain the heap invariant as necessary by repeatedly
+--- swapping with the smallest of the two children.
+---@param i (number) the starting position; default to the root.
+function BinaryHeap:trickle_down(i)
+  i = i or 1
+  local left, a, comp = i << 1, self, self.comp
+  while left <= #a do
+    local right = left + 1
+    local right = right <= #a and right or left -- 'right' is out of bound
+    local child = comp(a[left], a[right]) and left or right
+    if comp(a[i], a[child]) then
+      return
+    end
+    a[i], a[child] = a[child], a[i]
+    i = child
+    left = i << 1
+  end
+end
+
+--- Move the last element to the root, and then maintain the heap invariant as necessary by repeatedly
+--- swapping with the smallest of the two children.
+function BinaryHeap:remove()
+  local a = self
+  local root = a[1]
+  if #a == 1 then
+    a[1] = nil
+  else
+    a[1], a[#a] = a[#a], nil
+    self:trickle_down()
+  end
+  return root
+end
+
+--- Append the given value to the internal array, and then maintain the heap invariant by
+--- repeatedly swapping with the parent if necessary.
+---@param x any the value to be added to the heap
+function BinaryHeap:add(x)
+  local a = self
+  a[#a + 1] = x
+  self:bubble_up()
+  return a
+end
+
+function BinaryHeap:class(t, comp)
+  t = t or {}
+  local mt = { -- used to store the heap instance specific comparision function
+    comp = comp or function(a, b) -- default to min heap
+      return a <= b
+    end
+  }
+  setmetatable(mt, self)
+  self.__index = self
+
+  setmetatable(t, mt)
+  mt.__index = mt
+  return t
+end
+
+--- Build a heap from the given array.
+--- Starting from the lowest level and moving upwards, sift the root of each subtree downward
+--- as in the deletion algorithm until the heap property is restored.
+--- Courtesy of https://en.wikipedia.org/wiki/Robert_W._Floyd
+---@param a (table) an array of elements in arbitrary order
+function BinaryHeap:heapify()
+  local a = self
+  local parent = #a >> 1
+  local left = parent << 1
+  while parent > 0 do
+    self:trickle_down(parent)
+    left = left - 2
+    parent = left >> 1
+  end
+  return a
+end
+
+function BinaryHeap:new(a, comp)
+  a = a or {}
+  return self:class(a, comp):heapify()
+end
+
+function BinaryHeap:newMaxHeap(a) -- a convenient method to build a max heap
+  return self:new(a, function(x, y)
+    return x >= y
+  end)
+end
+
+function BinaryHeap:size()
+  return #self
+end
+
+function BinaryHeap:dump() -- debugging
+  return table.concat(self, ",")
+end
+
+function BinaryHeap:verify() -- debugging: verify the heap invariant holds
+  local a, comp = self, self.comp
+  for i = 1, #a do
+    local left = i << 1
+    if left > #a then
+      return
+    end
+    local right = left + 1
+    right = right > #a and left or right
+    local child = comp(a[left], a[right]) and left or right
+    assert(comp(a[i], a[child]))
+  end
+end
+return BinaryHeap