heap.rb in compsci-0.0.3.1

- old
+ new
@@ -1,28 +1,137 @@
 require 'compsci/tree'
 
-# A Heap is a partially sorted, complete binary tree with the property:
-# * Every node has a value larger (or smaller) than that of its children.
+include CompSci
+
+# A Heap is a partially sorted, complete N-ary tree with the property:
+# * Every node has a value larger (or smaller) than that of its children
+#   (the heap property is satisfied when a parent value equals a child value)
 #
-# This class implements a heap using a simple array for storage.
-# Array index math is used to find:
-# * The root node (idx 0)
-# * The "bottom-most" leaf node (last idx)
-# * Parent idx (idx-1 / 2)
-# * Child idx (2*idx + 1, 2*idx + 2)
-#
-# Any Comparable may be used for node values.
-# Initialize a heap with a cmp_val, either 1 for a MaxHeap or -1 for a MinHeap.
-# The heap property is satisfied when a parent value equals a child value.
-# Insertion (push) and removal (pop) are O(log n) where n is the heap size.
-# Nodes are inserted at the end of the array, and sift_up is called to
+# Implementation details:
+# * Any Comparable may be used for node values.
+# * Initialize a heap with a cmp_val, either 1 for a MaxHeap or -1 for a
+#   MinHeap.
+# * Insertion (push) and removal (pop) are O(logb n) where n is the heap size
+#   and b is child_slots (the N in N-ary)
+# * Nodes are inserted at the end of the array, and sift_up is called to
 #   reestablish the heap property.
-# Nodes are removed from the start of the array, and sift_down is called to
+# * Nodes are removed from the start of the array, and sift_down is called to
 #   reestablish the heap property.
-# Sift_up and sift_down are O(log n) because they only have to check and swap
-#   nodes at each layer of the tree, and there are log n layers to the tree.
+# * Sift_up and sift_down are O(logb n) because they only have to check and
+#   swap nodes at each layer of the tree, and there are log(n, base b) layers
+#   to the tree.
 #
-class CompSci::Heap < CompSci::CompleteBinaryTree
+class Heap < CompleteNaryTree
+  # * defaults to a MaxHeap, with the largest node at the root
+  # * specify a minheap with minheap: true or cmp_val: -1
+  #
+  def initialize(cmp_val: 1, minheap: false, child_slots: 2)
+    super(child_slots: child_slots)
+    cmp_val = -1 if minheap
+    case cmp_val
+    when -1, 1
+      @cmp_val = cmp_val
+    else
+      raise(ArgumentError, "unknown comparison value: #{cmp_val}")
+    end
+  end
+
+  # * append to the array
+  # * sift_up -- O(log n) on heap size
+  #
+  def push(node)
+    @store << node
+    self.sift_up(@store.size - 1)
+  end
+
+  # * remove from the front of the array
+  # * move last node to root
+  # * sift_down -- O(log n) on heap size
+  #
+  def pop
+    node = @store.shift
+    replacement = @store.pop
+    @store.unshift replacement if replacement
+    self.sift_down(0)
+    node
+  end
+
+  # * return what pop would return (avoid sifting)
+  #
+  def peek
+    @store.first
+  end
+
+  # * called recursively
+  # * idx represents the node suspected to violate the heap
+  # * intended to be O(log n) on heap size (log base child_slots)
+  #
+  def sift_up(idx)
+    return self if idx <= 0
+    pidx = self.class.parent_idx(idx, @child_slots)
+    if !self.heapish?(pidx, idx)
+      @store[idx], @store[pidx] = @store[pidx], @store[idx]   # swap
+      self.sift_up(pidx)
+    end
+    self
+  end
+
+  # * called recursively
+  # * idx represents the node suspected to violate the heap
+  # * intended to be O(log n) on heap size (log base child_slots)
+  # * slower than sift_up because one parent vs multiple children
+  #
+  def sift_down(idx)
+    return self if idx >= @store.size
+    cidxs = self.class.children_idx(idx, @child_slots)
+    # promote the heapiest child
+    cidx = self.heapiest(cidxs)
+    if !self.heapish?(idx, cidx)
+      @store[idx], @store[cidx] = @store[cidx], @store[idx]   # swap
+      self.sift_down(cidx)
+    end
+    self
+  end
+
+  # are values of parent and child (by index) in accordance with heap property?
+  #
+  def heapish?(pidx, cidx)
+    (@store[pidx] <=> @store[cidx]) != (@cmp_val * -1)
+  end
+
+  # return the heapiest idx in cidxs
+  #
+  def heapiest(cidxs)
+    idx = cidxs.first
+    cidxs.each { |cidx|
+      idx = cidx if cidx < @store.size and self.heapish?(cidx, idx)
+    }
+    idx
+  end
+
+  # * not used internally
+  # * checks that every node satisfies the heap property
+  # * calls heapish? on idx's children and then heap? on them recursively
+  #
+  def heap?(idx: 0)
+    check_children = []
+    self.class.children_idx(idx, @child_slots).each { |cidx|
+      # cidx is arithmetically produced; the corresponding child may not exist
+      if cidx < @store.size
+        return false unless self.heapish?(idx, cidx)
+        check_children << cidx
+      end
+    }
+    check_children.each { |cidx| return false unless self.heap?(idx: cidx) }
+    true
+  end
+end
+
+#
+# LEGACY BELOW untouched for now
+#
+
+class BinaryHeap < CompleteBinaryTree
   # defaults to a MaxHeap, with the largest node at the root
   # specify a minheap with minheap: true or cmp_val: -1
   #
   def initialize(cmp_val: 1, minheap: false)
     super()