ハフマンツリーをRubyで書く

ハフマン符号化に使うハフマンツリーをRubyで作ってみた。2本のキューを用意して、2つのノードをくっつけたり出し入れするのが簡単というのでそういう方針で。本当は1本はプライオリティ・キューでないとエンキュー時にソートするために効率が悪い。

# Huffman coding
# 

class Node
  attr_accessor :val, :weight, :left, :right

  def initialize(val = "", weight = 0)
    @val, @weight = val, weight
  end

  def leaf?
    !(left and right)
  end

end

class HuffmanTree
  attr_accessor :freq, :que1, :que2, :root

  def initialize(str)
    @freq = count_freq(str.downcase)
    @que1 = build_queue(@freq)
    @que2 = []
    build
  end

  def build
    while @que1.size + @que2.size > 1
      left  = deque_the_lower(@que1, @que2)
      right = deque_the_lower(@que1, @que2)
      node = Node.new(left.val + right.val, left.weight + right.weight)
      node.left, node.right = left, right
      enque(@que2, node)
    end
    @root = @que2
  end

  def deque_the_lower(q1, q2)
    return q1.shift if q2.size == 0
    return q2.shift if q1.size == 0
    if q1.first.weight < q2.first.weight
      q1.shift
    else
      q2.shift
    end
  end

  def enque(q, node)
    q << node
    q.sort_by{|x| -x.weight}  # should use priority que insted
  end 

  def count_freq(str)
    str.split(//).inject(Hash.new(0)) do |hash, char|
      hash[char] += 1
      hash
    end
  end

  def build_queue(freq)
    freq.to_a.sort_by{|x| x[1]}.map{|f| Node.new(f[0], f[1])}
  end
end

if __FILE__ == $0
  h = HuffmanTree.new("this is an example of a huffman tree")
  p h.root
end

Graphvizでビジュアル化してみた。

require './huff.rb'
require 'graphviz'

def create_node(g, n, prev = nil)
  if n.leaf?
    new_node = g.add_node("\'#{n.val}\' (#{n.weight.to_s})", :color => :red)
    g.add_edge(prev, new_node) if prev
    return
  else
    new_node = g.add_node("\'#{n.val}\' (#{n.weight.to_s})", :color => :blue)
    g.add_edge(prev, new_node) if prev
    create_node(g, n.left , new_node) if n.left
    create_node(g, n.right, new_node) if n.right
  end
end

h = HuffmanTree.new("this is an example of a huffman tree")

g = GraphViz::new("G")
create_node(g, h.root[0])
g.output(:png => "graph.png")