Instantly share code, notes, and snippets.

# wmwong/heap.rb Created Sep 28, 2013

Fun with heaps
 module Heap def parent_index(i) # +1 because the math works with indices starting from 1. # -1 because we need to convert back to 0-index. (i + 1) / 2 - 1 end def left_child_index(i) # +1 because the math works with indices starting from 1. # -1 because we need to convert back to 0-index. # (i + 1) * 2 - 1 # Use the simplified form. i * 2 + 1 end def right_child_index(i) left_child_index(i) + 1 end # Sort the array according to the order. # order can either be :inc or :dec. def sort(array, order = :inc) order == :inc ? sort_increasing(array) : sort_decreasing(array) end # Sort the array in increasing order. def sort_increasing(array) create_max_heap(array) (array.length - 1).downto(1).each do |i| temp = array[i] array[i] = array[0] array[0] = temp max_heapify(array, 0, i) end end # Create a max-heap from the given array. def create_max_heap(array) return unless array.length > 1 # -1 because 0-index. ((array.length / 2) - 1).downto(0).each do |i| max_heapify(array, i, array.length) end end # The left and right subtrees of the node at the given index must already be # max-heaps. def max_heapify(array, index, heap_length) return if array.empty? # Find the larger of parent, left child, and right child. left_index = left_child_index(index) right_index = right_child_index(index) largest_index = index largest_value = array[index] if left_index < heap_length && array[left_index] > array[index] largest_index = left_index largest_value = array[left_index] end if right_index < heap_length && array[right_index] > largest_value largest_index = right_index end # If the largest is a child, swap the parent with the child. unless largest_index == index temp = array[largest_index] array[largest_index] = array[index] array[index] = temp max_heapify(array, largest_index, heap_length) end end # Sort the array in decreasing order. def sort_decreasing(array) create_min_heap(array) (array.length - 1).downto(1).each do |i| temp = array[i] array[i] = array[0] array[0] = temp min_heapify(array, 0, i) end end # Create a min-heap from the given array. def create_min_heap(array) return unless array.length > 1 ((array.length / 2) - 1).downto(0).each do |i| min_heapify(array, i, array.length) end end # The left and right subtrees of the node at the given index must already be # min-heaps. def min_heapify(array, index, heap_length) return if array.empty? # Find the smallest of parent, left child, and right child. left_index = left_child_index(index) right_index = right_child_index(index) smallest_index = index smallest_value = array[index] if left_index < heap_length && array[left_index] < array[index] smallest_index = left_index smallest_value = array[left_index] end if right_index < heap_length && array[right_index] < smallest_value smallest_index = right_index end # If the smallest is a child, swap the parent with the child. unless smallest_index == index temp = array[smallest_index] array[smallest_index] = array[index] array[index] = temp min_heapify(array, smallest_index, heap_length) end end end
 require './heap' require 'test/unit' class HeapTest < Test::Unit::TestCase include Heap def test_parent_index assert_equal(2, parent_index(6)) assert_equal(2, parent_index(5)) assert_equal(1, parent_index(4)) assert_equal(1, parent_index(3)) assert_equal(0, parent_index(2)) assert_equal(0, parent_index(1)) end def test_left_child_index assert_equal(1, left_child_index(0)) assert_equal(3, left_child_index(1)) assert_equal(5, left_child_index(2)) end def test_right_child_index assert_equal(2, right_child_index(0)) assert_equal(4, right_child_index(1)) assert_equal(6, right_child_index(2)) end def test_max_heapify # Empty array. array = [] max_heapify(array, 0, 0) assert(array.empty?) # Single element. array = [1] max_heapify(array, 0, 1) assert_equal(1, array[0]) # Two levels. array = [1, 2] max_heapify(array, 0, 2) assert_equal(2, array[0]) assert_equal(1, array[1]) # Two levels, full. array = [1, 2, 3] max_heapify(array, 0, 3) assert_equal(3, array[0]) assert_equal(2, array[1]) assert_equal(1, array[2]) # Three levels. array = [1, 4, 3, 2] max_heapify(array, 0, 4) assert_equal(4, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) assert_equal(1, array[3]) # Respects heap length. array = [1, 4, 3, 2, 5, 6, 7] max_heapify(array, 0, 4) assert_equal(4, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) assert_equal(1, array[3]) assert_equal(5, array[4]) assert_equal(6, array[5]) assert_equal(7, array[6]) end def test_create_max_heap # Empty array. array = [] create_max_heap(array) assert(array.empty?) # Single element. array = [1] create_max_heap(array) assert_equal(1, array[0]) # Two levels. array = [1, 2] create_max_heap(array) assert_equal(2, array[0]) assert_equal(1, array[1]) # Two levels, full. array = [1, 2, 3] create_max_heap(array) assert_equal(3, array[0]) assert_equal(2, array[1]) assert_equal(1, array[2]) # Three levels. array = [1, 2, 3, 4] create_max_heap(array) assert_equal(4, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) assert_equal(1, array[3]) end def test_sort_increasing # Empty array. array = [] sort_increasing(array) assert(array.empty?) # Single element. array = [1] sort_increasing(array) assert_equal(1, array[0]) # Two levels. array = [2, 1] sort_increasing(array) assert_equal(1, array[0]) assert_equal(2, array[1]) # Two levels, full. array = [3, 2, 1] sort_increasing(array) assert_equal(1, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) # Three levels. array = [4, 3, 2, 1] sort_increasing(array) assert_equal(1, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) assert_equal(4, array[3]) end def test_min_heapify # Empty array. array = [] min_heapify(array, 0, 0) assert(array.empty?) # Single element. array = [1] min_heapify(array, 0, 1) assert_equal(1, array[0]) # Two levels. array = [2, 1] min_heapify(array, 0, 2) assert_equal(1, array[0]) assert_equal(2, array[1]) # Two levels, full. array = [3, 2, 1] min_heapify(array, 0, 3) assert_equal(1, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) # Three levels. array = [4, 1, 3, 2] min_heapify(array, 0, 4) assert_equal(1, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) assert_equal(4, array[3]) # Respects heap length. array = [4, 1, 3, 2, 5, 6, 7] min_heapify(array, 0, 4) assert_equal(1, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) assert_equal(4, array[3]) assert_equal(5, array[4]) assert_equal(6, array[5]) assert_equal(7, array[6]) end def test_create_min_heap # Empty array. array = [] create_min_heap(array) assert(array.empty?) # Single element. array = [1] create_min_heap(array) assert_equal(1, array[0]) # Two levels. array = [2, 1] create_min_heap(array) assert_equal(1, array[0]) assert_equal(2, array[1]) # Two levels, full. array = [3, 2, 1] create_min_heap(array) assert_equal(1, array[0]) assert_equal(2, array[1]) assert_equal(3, array[2]) # Three levels. array = [4, 3, 2, 1] create_min_heap(array) assert_equal(1, array[0]) assert_equal(3, array[1]) assert_equal(2, array[2]) assert_equal(4, array[3]) end def test_sort_decreasing # Empty array. array = [] sort_decreasing(array) assert(array.empty?) # Single element. array = [1] sort_decreasing(array) assert_equal(1, array[0]) # Two levels. array = [1, 2] sort_decreasing(array) assert_equal(2, array[0]) assert_equal(1, array[1]) # Two levels, full. array = [1, 2, 3] sort_decreasing(array) assert_equal(3, array[0]) assert_equal(2, array[1]) assert_equal(1, array[2]) # Three levels. array = [1, 2, 3, 4] sort_decreasing(array) assert_equal(4, array[0]) assert_equal(3, array[1]) assert_equal(2, array[2]) assert_equal(1, array[3]) end end
 require './heap' module PriorityQueue class Max include Heap # For testing purposes. attr_reader :array def initialize(array) @array = array create_max_heap(@array) end def get_max @array[0] end def extract_max return nil if @array.empty? max = get_max temp = @array.pop unless @array.empty? @array[0] = temp max_heapify(@array, 0, @array.length) end max end def increase(index, obj) raise "#{obj} is less than #{@array[index]} and is not an increase." if obj < @array[index] @array[index] = obj while index > 0 && @array[parent_index(index)] < obj temp = @array[index] @array[index] = @array[parent_index(index)] @array[parent_index(index)] = temp index = parent_index(index) end end def insert(obj) @array.push(obj) increase(@array.length - 1, obj) end end class Min include Heap # For testing purposes. attr_reader :array def initialize(array) @array = array create_min_heap(@array) end def get_min @array[0] end def extract_min return nil if @array.empty? min = get_min temp = @array.pop unless @array.empty? @array[0] = temp min_heapify(@array, 0, @array.length) end min end def decrease(index, obj) raise "#{obj} is greater than #{@array[index]} and is not a decrease." if obj > @array[index] @array[index] = obj while index > 0 && @array[parent_index(index)] > obj temp = @array[index] @array[index] = @array[parent_index(index)] @array[parent_index(index)] = temp index = parent_index(index) end end def insert(obj) @array.push(obj) decrease(@array.length - 1, obj) end end end
 require './priority_queue' require 'test/unit' class TestMaxPriorityQueue < Test::Unit::TestCase def test_get_max queue = PriorityQueue::Max.new([1, 2, 3, 4, 5, 6]) assert_equal(6, queue.get_max) end def test_extract_max queue = PriorityQueue::Max.new([1, 2, 3, 4, 5, 6]) assert_equal(6, queue.extract_max) assert_equal(5, queue.extract_max) assert_equal(4, queue.extract_max) assert_equal(3, queue.extract_max) assert_equal(2, queue.extract_max) assert_equal(1, queue.extract_max) assert_equal(nil, queue.extract_max) end def test_increase queue = PriorityQueue::Max.new([1, 2, 3, 4, 5, 6]) assert_raise RuntimeError do queue.increase(0, 0) end queue.increase(5, 7) assert_equal(7, queue.array[0]) assert_equal(5, queue.array[1]) assert_equal(6, queue.array[2]) assert_equal(4, queue.array[3]) assert_equal(2, queue.array[4]) assert_equal(3, queue.array[5]) end def test_insert queue = PriorityQueue::Max.new([1, 2, 3, 4, 5, 6]) queue.insert(7) assert_equal(7, queue.array[0]) assert_equal(5, queue.array[1]) assert_equal(6, queue.array[2]) assert_equal(4, queue.array[3]) assert_equal(2, queue.array[4]) assert_equal(1, queue.array[5]) assert_equal(3, queue.array[6]) end end class TestMinPriorityQueue < Test::Unit::TestCase def test_get_min queue = PriorityQueue::Min.new([6, 5, 4, 3, 2, 1]) assert_equal(1, queue.get_min) end def test_extract_min queue = PriorityQueue::Min.new([6, 5, 4, 3, 2, 1]) assert_equal(1, queue.extract_min) assert_equal(2, queue.extract_min) assert_equal(3, queue.extract_min) assert_equal(4, queue.extract_min) assert_equal(5, queue.extract_min) assert_equal(6, queue.extract_min) assert_equal(nil, queue.extract_min) end def test_increase queue = PriorityQueue::Min.new([6, 5, 4, 3, 2, 1]) assert_raise RuntimeError do queue.decrease(0, 7) end queue.decrease(5, 0) assert_equal(0, queue.array[0]) assert_equal(2, queue.array[1]) assert_equal(1, queue.array[2]) assert_equal(3, queue.array[3]) assert_equal(5, queue.array[4]) assert_equal(4, queue.array[5]) end def test_insert queue = PriorityQueue::Min.new([6, 5, 4, 3, 2, 1]) queue.insert(0) assert_equal(0, queue.array[0]) assert_equal(2, queue.array[1]) assert_equal(1, queue.array[2]) assert_equal(3, queue.array[3]) assert_equal(5, queue.array[4]) assert_equal(6, queue.array[5]) assert_equal(4, queue.array[6]) end end