LZWをRubyで実装してみた

lzw.rb

#
# LZW Algorithm
#
class Dic
  def initialize
    @dict = (0...256).map{|k| [nil, k]}
    @code = {}
  end

  def add(w, k)
    new_code = @code[(w << 8) + k] = @dict.size
    @dict << [w, k]
    return new_code
  end

  def code(w, k)
    return k unless w
    return @code[(w << 8) + k]
  end

  def string(code)
    w, k = @dict[code]
    return nil unless k
    return string(w) << k if w
    return [k]
  end
end

def encode(input)
  output = []
  dic = Dic.new
  w = nil
  while k = input.shift
    c = dic.code(w, k)
    if c
      w = c
    else
      output << w
      dic.add(w, k)
      w = k
    end
  end
  output << w if w
  return output
end

def decode(input)
  output = []
  dic = Dic.new
  w = input.shift
  return output unless w
  ws = dic.string(w)
  ws.each{|v| output << v}
  while k = input.shift
    ks = dic.string(k)
    ks = ws + ws.slice(0, 1) unless ks
    ks.each{|v| output << v}
    dic.add(w, ks[0])
    w, ws = k, ks
  end
  return output
end

encode([0,1,0,1,0,1,0,1]) # => [0, 1, 256, 258, 1]
decode([0,1,256,258,1])   # => [0, 1, 0, 1, 0, 1, 0, 1]