Erlangとハフマン符号

他の言語で書いたことがあるコードをErlangで書くとどうなるかいろいろ試してみたくなりました。第1弾はハフマン符号です。

ハフマン符号化の大まかな流れは、(1) 頻度表を作る、(2) 頻度表を元にハフマン木を作る、(3) ハフマン木を元に入力記号→ビット列の対応表を作る、(4) 3で作った対応表を元に符号化する、となります。

ではさっそく作ってみましょう。まずは「(1) 頻度表を作る」です。

make_frqtbl(Bin) ->
    Tmp =
	lists:foldl(
	  fun(B, Ary) ->
		  array:set(B, array:get(B, Ary) + 1, Ary)
	  end, array:new(256 + 1, {default, 0}), binary_to_list(Bin)),
    {Result, _} = 
	lists:foldl(
	  fun(X, {Res, I}) ->
		  {[{I, X}|Res], I + 1}
	  end, {[], 0}, array:to_list(Tmp)),
    lists:filter(fun({_,C}) -> C > 0 end, lists:reverse(Result)).

以前勉強した array を使って頻度をカウントした後、 {符号, 頻度} のタプルからなるリストを作っています。

7> Input = <<"ABCABA">>.
<<"ABCABA">>
8> Frq = huffman:make_frqtbl(Input).
[{65,3},{66,2},{67,1}]

次は「(2) 頻度表を元にハフマン木を作る」です。

make_huffman_tree([Node]) -> Node;
make_huffman_tree(Tbl) ->
    STbl = lists:sort(fun({_,C1},{_,C2}) -> C1 < C2 end, Tbl),
    [{X1,C1},{X2,C2}|TL] = STbl,
    make_huffman_tree(
      [
       {[{X1,C1},{X2,C2}], C1+C2} | TL
      ]).

頻度の小さいもの2つを一つのノードにまとめて...というのをノードが一つになるまで繰り返します。まとめたノードを適切な場所に移動させるために毎回ソートさせるという頭の悪いやりかたをしています。

9> Tree = huffman:make_huffman_tree(Frq).
{[{65,3},{[{67,1},{66,2}],3}],6}

見づらいですが、ルートから左に1つ降りるとA、右に降りるとBとCの枝があります。左に辿るときに0b、右に辿るときに1bというビットを割り振って、各符号がどのような経路で辿りつけるかを調べて対応表を作ります。これが「(3) ハフマン木を元に入力記号→ビット列の対応表を作る」です。

get_encodelist(Tbl) ->
    lists:flatten(get_encodelist(Tbl, <<>>)).
get_encodelist([], _) -> [];
get_encodelist({X,_C}, Code) -> 
    case X of 
	X when is_integer(X) ->
	    {X, Code};
	[L,R] ->
	    [get_encodelist(L, <<Code/bitstring, 0:1>>),
	    get_encodelist(R, <<Code/bitstring, 1:1>>)]
    end.    

bitstring というのは今回初めて存在を知りました。バイナリデータだと /binary という指定で連結できますが、 bitstringはビット列の連結に使えるようです。

11> Sym = huffman:get_encodelist(Tree).
[{65,<<0:1>>},{67,<<2:2>>},{66,<<3:2>>}]

よいようです。あとはこれらの情報を元に1バイトずつ符号化してゆけば完成です。(「(4) 3で作った対応表を元に符号化する」)

encode(Bin) ->
    Frq = make_frqtbl(Bin),
    Tree = make_huffman_tree(Frq),
    Sym = get_encodelist(Tree),
    Dic = 
	lists:foldl(
	  fun({X, Bits}, Acc) ->
		  dict:append(X, Bits, Acc)
	  end, dict:new(), Sym),
    {Tree,
     lists:foldl(
       fun(B, Acc) ->
	       {ok, [D]} = dict:find(B, Dic),
	       <<Acc/bitstring, D/bitstring>>
       end, <<>>, binary_to_list(Bin))}.

対応表をいったん dict に構成しなおして、それを1バイトずつやっつけているだけですね。バイト境界を気にせず入出力できるのが楽ちんです。

13> {Tree, Bin} = huffman:encode(<<"ABCABA">>).
{{[{65,3},{[{67,1},{66,2}],3}],6},<<115,0:1>>}
14> bit_size(Bin).
9

元のバイト列が48ビットだったのに対し、符号化後は9ビットになっています。

次は復号です。復号には、符号化列のほかにハフマン木が必要となります。ハフマン木は頻度表からも作ることができますので、頻度表か、ハフマン木のいずれかが残っていればOKです。

decode({Tree, Bin}) -> decode(Tree, Bin).
decode(Tree, Bin) ->
    decode(Tree, Bin, <<>>).

decode(_Tree, <<>>, Res) -> Res;
decode(Tree, Bin, Res) -> 
    {B, Rem} = decode_core(Tree, Bin),
    decode(Tree, Rem, <<Res/binary, B/binary>>).

decode_core({X,_}, Bin) when is_integer(X) -> {<<X>>, Bin};
decode_core({[L,R],_}, <<B:1, Rem/bitstring>>) -> 
    case B of
	0 ->
	    decode_core(L, Rem);
	1 ->
	    decode_core(R, Rem)
    end.

復号は1ビットずつ読んでいってハフマン木を辿り、枝に行きついたらそれを出力するという繰り返しになります。

15> huffman:decode(Tree, Bin).
<<"ABCABA">>

よいようです。もうちょっとでかいデータを圧縮してみましょう。

19> {ok, Input} = file:read_file("huffman.erl").
{ok,<<"-module(huffman).\r\n-compile(export_all).\r\n\r\nmake_frqtbl(Bin) ->\r\n    Tmp =\r\n\tlists:foldl(\r\n\t  fun(B, Ary) ->"...>>}
20> bit_size(Input).
14784
21> {Tree, Compressed} = huffman:encode(Input).
{{[{[{[{[{[{99,41},{[{84,21},{109,22}],43}],84},{44,86}],
        170},
(中略)
  1848},
 <<84,51,66,81,45,175,114,159,124,57,34,227,153,210,160,
   48,103,115,37,182,235,1,217,153,77,164,33,...>>}
22> bit_size(Compressed).
9421

14784ビットが9421ビットに縮まりました。(しつこいですが復号のためには9421ビットの符号化列のほかにハフマン木も必要です)

さて、ちゃんと戻るでしょうか...

23> Output = huffman:decode(Tree, Compressed).
<<"-module(huffman).\r\n-compile(export_all).\r\n\r\nmake_frqtbl(Bin) ->\r\n    Tmp =\r\n\tlists:foldl(\r\n\t  fun(B, Ary) ->\r\n\t\t  ar"...>>
24> Input = Output.
<<"-module(huffman).\r\n-compile(export_all).\r\n\r\nmake_frqtbl(Bin) ->\r\n    Tmp =\r\n\tlists:foldl(\r\n\t  fun(B, Ary) ->\r\n\t\t  ar"...>>
25> 

OKのようです。

性能無視で書いているのでこのままでは実用に耐えませんが(そもそもErlang nativeで書く意味もあんまりないように思いますが)、配列まわり以外は比較的素直に書けました。バイナリ処理が自然に書けるというErlangの利点を生かして本来の処理に集中できたように思います。

最後に、ソース全体を載せておきます。

-module(huffman).
-compile(export_all).

make_frqtbl(Bin) ->
    Tmp =
	lists:foldl(
	  fun(B, Ary) ->
		  array:set(B, array:get(B, Ary) + 1, Ary)
	  end, array:new(256 + 1, {default, 0}), binary_to_list(Bin)),
    {Result, _} = 
	lists:foldl(
	  fun(X, {Res, I}) ->
		  {[{I, X}|Res], I + 1}
	  end, {[], 0}, array:to_list(Tmp)),
    lists:filter(fun({_,C}) -> C > 0 end, lists:reverse(Result)).

make_huffman_tree([Node]) -> Node;
make_huffman_tree(Tbl) ->
    STbl = lists:sort(fun({_,C1},{_,C2}) -> C1 < C2 end, Tbl),
    [{X1,C1},{X2,C2}|TL] = STbl,
    make_huffman_tree(
      [
       {[{X1,C1},{X2,C2}], C1+C2} | TL
      ]).

get_encodelist(Tbl) ->
    lists:flatten(get_encodelist(Tbl, <<>>)).
get_encodelist([], _) -> [];
get_encodelist({X,_C}, Code) -> 
    case X of 
	X when is_integer(X) ->
	    {X, Code};
	[L,R] ->
	    [get_encodelist(L, <<Code/bitstring, 0:1>>),
	    get_encodelist(R, <<Code/bitstring, 1:1>>)]
    end.    
    
encode(Bin) ->
    Frq = make_frqtbl(Bin),
    Tree = make_huffman_tree(Frq),
    Sym = get_encodelist(Tree),
    Dic = 
	lists:foldl(
	  fun({X, Bits}, Acc) ->
		  dict:append(X, Bits, Acc)
	  end, dict:new(), Sym),
    {Tree,
     lists:foldl(
       fun(B, Acc) ->
	       {ok, [D]} = dict:find(B, Dic),
	       <<Acc/bitstring, D/bitstring>>
       end, <<>>, binary_to_list(Bin))}.

decode({Tree, Bin}) -> decode(Tree, Bin).
decode(Tree, Bin) ->
    decode(Tree, Bin, <<>>).

decode(_Tree, <<>>, Res) -> Res;
decode(Tree, Bin, Res) -> 
    {B, Rem} = decode_core(Tree, Bin),
    decode(Tree, Rem, <<Res/binary, B/binary>>).

decode_core({X,_}, Bin) when is_integer(X) -> {<<X>>, Bin};
decode_core({[L,R],_}, <<B:1, Rem/bitstring>>) -> 
    case B of
	0 ->
	    decode_core(L, Rem);
	1 ->
	    decode_core(R, Rem)
    end.