Mae向きなブログ

Mae向きな情報発信を続けていきたいと思います。

「リンゴ列をもっと短く！」をGvizで描く

Ruby CodeIQ

予てからGvizを使って、ビジュアル化をやってみたいと思っていたのですが、なかなか取り組めずにいました。何か良い題材はないかと思っていたのですが、先日、解いた「リンゴ列をもっと短く！」のハフマン木を実際に見てみたいと思い取り組んでみました。

apple_gviz.rb

# -*- coding: utf-8 -*-
require 'gviz'

class Node
  attr_accessor :freq, :left, :center, :right
  def initialize(freq, left = nil, center = nil, right = nil)
    @freq = freq
    @left, @center, @right = left, center, right
  end
end

class Leaf < Node
  attr_accessor :ch
  def initialize(ch, freq)
    @ch, @freq = ch, freq
  end
end

typical_text = File.open('typical.txt').read
freq = Hash.new(0)
# 各単語の出現頻度を調べる
typical_text.each_char {|ch|
  freq[ch] += 1
}

freq["dummy"] = 0
arr = []
freq.each {|k, v|
  arr << Leaf.new(k, v)
}
loop {
  break if arr.size == 1
  arr.sort_by!{|e| e.freq }
  if arr.size == 2
    l = arr.shift
    c = arr.shift
    arr << Node.new((l.freq+c.freq), l, c)    
  else
    l = arr.shift
    c = arr.shift
    r = arr.shift
    arr << Node.new((l.freq+c.freq+r.freq), l, c, r)
  end
}

$path = ""              
$ans = Hash.new
$gviz_hash = Hash.new

def traverse parent, node, code = nil
  $path << code unless code.nil?
  $gviz_hash[parent.object_id.to_s] = [] if $gviz_hash[parent.object_id.to_s].nil?
  if node.instance_of?(Leaf)
    $gviz_hash[parent.object_id.to_s] << node.ch
  else
    $gviz_hash[parent.object_id.to_s] << node.object_id.to_s
  end
  if node.instance_of?(Leaf)
    $ans[node.ch] = $path.dup
    $path.chop             
    return
  end
  traverse node, node.left, "r"  unless node.left.nil?
  traverse node, node.center, "g" unless node.center.nil?
  traverse node, node.right, "b"  unless node.right.nil?
  $path.chop
end

root = arr.first 
traverse root, root

gv = Gviz.new
gv.nodes shape:'circle'
$gviz_hash.each {|key, values|
  values.each {|node|
    next if key == node
    gv.route(key => node)
    gv.node key.to_sym, label: "" unless [*'A'..'Z'].include?(key)
    gv.node node.to_sym, label: "_" if node == "dummy"
  }
}
gv.save("apple_tree", :png)

実行結果

以下が実行結果です。@merborne さんのように綺麗に描けませんが、少しずつ勉強していこうと思います。

参考