二分ヒープの Ruby 実装

max_heap.rb

Node = Struct.new(:key, :value)

class MaxHeap
  def initialize
    @heap = [nil]
    @size = 0
  end
  attr_reader :heap, :size
  
  def insert(key, value)
    @heap << Node.new(key, value)
    @size += 1
    up_heap(@size / 2)
  end
  alias :[]= :insert
  
  def up_heap(i)
    return if i.zero?
    l, r = i * 2, i * 2 + 1
    largest = (l <= @size and @heap[l].key > @heap[i].key) ? l : i
    largest = r if r <= @size and @heap[r].key > @heap[largest].key
    unless largest == i
      @heap[i], @heap[largest] = @heap[largest], @heap[i]
      up_heap(i / 2)
    end
  end
   
  def remove_root
    return nil if @size.zero?
    a = @heap[1]
    b = @heap.pop
    @size -= 1
    return b if @size.zero?
    @heap[1] = b
    down_heap(1)
    a
  end
   
  def down_heap(i)
    return if i > @size
    l, r = i * 2, i * 2 + 1
    largest = (l <= @size and @heap[l].key > @heap[i].key) ? l : i
    largest = r if r <= @size and @heap[r].key > @heap[largest].key
    unless largest == i
      @heap[i], @heap[largest] = @heap[largest], @heap[i]
      down_heap(largest)
    end
  end
  
  # Nodeを返す
  def search(key)
    @heap[1..-1].find {|node| node.key == key}
  end
  
  # valueを返す
  def [](key)
    search(key)&.value
  end
  
  # rootを取り出す
  def root
    [@heap[1].key, @heap[1].value] if @heap[1]
  end
  
  def each
    @heap[1..-1].each {|node| yield(node.key, node.value)}
  end
  
  def print
    s = @heap.size
    traverse = ->(i, len = 0) {
      st = " -- " + @heap[i].key.to_s
      l = st.length
      return st if i + 1 >= s
      if i * 2 < s
        st += traverse.(i * 2, len + l)
        if i * 2 + 1 < s
          st += "\n" + " " * (len + l) + traverse.(i * 2 + 1, len + l)
        end
      end
      st 
    }
    puts traverse.(1)
  end
  
  include Enumerable
end

min_heap.rb

Node = Struct.new(:key, :value)

class MinHeap
  def initialize
    @heap = [nil]
    @size = 0
  end
  attr_reader :heap, :size
  
  def insert(key, value)
    @heap << Node.new(key, value)
    @size += 1
    up_heap(@size / 2)
  end
  alias :[]= :insert
  
  def up_heap(i)
    return if i.zero?
    l, r = i * 2, i * 2 + 1
    smallest = (r <= @size and @heap[r].key < @heap[i].key) ? r : i
    smallest = l if l <= @size and @heap[l].key < @heap[smallest].key
    unless smallest == i
      @heap[i], @heap[smallest] = @heap[smallest], @heap[i]
      up_heap(i / 2)
    end
  end
   
  def remove_root
    return nil if @size.zero?
    a = @heap[1]
    b = @heap.pop
    @size -= 1
    return b if @size.zero?
    @heap[1] = b
    down_heap(1)
    a
  end
   
  def down_heap(i)
    return if i > @size
    l, r = i * 2, i * 2 + 1
    smallest = (r <= @size and @heap[r].key < @heap[i].key) ? r : i
    smallest = l if l <= @size and @heap[l].key < @heap[smallest].key
    unless smallest == i
      @heap[i], @heap[smallest] = @heap[smallest], @heap[i]
      down_heap(smallest)
    end
  end
  
  # Nodeを返す
  def search(key)
    @heap[1..-1].find {|node| node.key == key}
  end
  
  # valueを返す
  def [](key)
    search(key)&.value
  end
  
  # rootを取り出す
  def root
    [@heap[1].key, @heap[1].value] if @heap[1]
  end
  
  def each
    @heap[1..-1].each {|node| yield(node.key, node.value)}
  end
  
  def print
    s = @heap.size
    traverse = ->(i, len = 0) {
      st = " -- " + @heap[i].key.to_s
      l = st.length
      return st if i + 1 >= s
      if i * 2 < s
        st += traverse.(i * 2, len + l)
        if i * 2 + 1 < s
          st += "\n" + " " * (len + l) + traverse.(i * 2 + 1, len + l)
        end
      end
      st 
    }
    puts traverse.(1)
  end
  
  include Enumerable
end