🐷

ハフマン符号をPythonで実装する

4 min read

はじめに

ハフマン符号をPythonで実装しました。
ハフマン符号はデータ中の各文字の出現頻度に応じて長さの異なるビット列を割り当てることでデータ全体の圧縮をはかるアルゴリズムです。

ハフマン符号の理解は以下の記事を参考にさせて頂きました。
Pythonでハフマン符号化を実装してみた - Wonder Code
上の記事にも実装例がありますが計算量を改善する余地があったので手を加えて実装しています。

本記事では実装の際にハマった点や注意した点などをまとめておこうと思います。
実装したリポジトリは以下です。

https://github.com/kawagh/HuffmanCoding

実装

符号化の準備であるハフマン木の構築、符号化、復号化に関して順に説明していき、最後にコード全体を示します。

符号化(encode)

ハフマン木の構築

ハフマン符号では符号化のための準備としてハフマン木を作ります。
ハフマン木を構成するノードとして出現頻度、対応する文字、左右の子ノードの情報を持つ以下のようなdataclassを定義しました。

なおclassにおいてメンバ変数に今定義しているクラスを型の注釈としてつける記法はpython3.8においてはfrom __future__ import annotationsをする必要があります。

from __future__ import annotations
from dataclasses import dataclass
from typing import Optional

@dataclass
class Node:
    cnt: int
    char: Optional[str]
    left: Optional[Node] = None
    right: Optional[Node] = None

    def __lt__(self, other: Node):
        return self.cnt < other.cnt

ハフマン木はこのノードを葉として出現頻度の少ない順にノード2つを取り出し出現頻度を足したノードを新たに追加するという操作を繰り返して根までボトムアップに構築します。
この構築では出現頻度が最小のノードを取り出す操作を頻繁に行う必要があるので優先度付きキューを使用しています。
優先度付きキューでは最小値の取り出しやキューへの追加が要素数をNとしてO(logN)で行えます。この優先度の比較のためにNodeでは__lt__()メソッドを定義しています。

ハフマン木を用いて符号化

ハフマン木を構築した後に符号化の作業に入ります。トップダウンに根から左右に移動していき葉に達したときにその移動順に対応(左が0,右が1など)したビット列がその葉の各文字に対する符号になります。
葉に達した時に復号のための辞書も作っています。
以下に示すコードの_rec()メソッドがこの操作に対応しています。

復号化(decode)

復号化はデータを前から読み出して辞書に一致する符号が出てきたらその値を読む作業になります。
符号が他の符号の接頭辞にならない性質により一意に復号することが出来ます。

実装コード

以上を踏まえた実装コードが以下のようになります。
encodeの部分に関する補足として出現頻度を数える際に各種文字に対してcount()メソッドを行うと文字種類数が増えた場合に低速になるのでCounterを利用しています。
また文字の種類が1つの場合や空文字列の場合は個別に扱っています。

import heapq
from collections import Counter

class HuffmanCoding:
    """
    Static Huffman Coding
    """

    def __init__(self):
        self.root = None
        self.encode_dict = {}
        self.decode_dict = {}

    def encode(self, raw_datas: str) -> str:
        if raw_datas == "":
            raise ValueError("cannot encode empty string")
        nodes: list[Node] = [Node(v, k) for k, v in Counter(raw_datas).items()]
        heapq.heapify(nodes)
        if len(nodes) == 1:
            self.root = heapq.heappop(nodes)
            self._rec(self.root, "0")
            return len(raw_datas) * "0"

        while len(nodes) >= 2:
            min_node1 = heapq.heappop(nodes)
            min_node2 = heapq.heappop(nodes)
            parent_node = Node(min_node1.cnt + min_node2.cnt, None, min_node1, min_node2)
            heapq.heappush(nodes, parent_node)
        self.root = heapq.heappop(nodes)
        self._rec(self.root, "")
        res = "".join([self.encode_dict[c] for c in raw_datas])
        return res

    def _rec(self, node: Node, s) -> None:
        if node.char:  # leaf
            self.encode_dict[node.char] = s
            self.decode_dict[s] = node.char
            return
        self._rec(node.left, s + "0")
        self._rec(node.right, s + "1")

    def decode(self, data) -> str:
        result = ""
        cur = ""
        for bit in data:
            cur += bit
            if cur in self.decode_dict:
                result += self.decode_dict[cur]
                cur = ""
        return result

復号化と符号化を実際に行う

実際に適当な文字列に関して符号化と復号化をした例が以下です。

def main():
    input_str = "DAEBCBACBBBC"
    hc = HuffmanCoding()
    encoded_str = hc.encode(input_str)
    print(f"{input_str=}")
    print(hc.encode_dict)
    print(f"{encoded_str=}")
    decoded_str = hc.decode(encoded_str)
    total_bitlength_iffixed = len(set(input_str)).bit_length() * len(input_str)
    rate = len(encoded_str) / total_bitlength_iffixed  # 固定長のbit列で文字を表現した場合とハフマン符号化した場合のbit長の比
    print(f"{rate=}")
    assert input_str == decoded_str


if __name__ == "__main__":
    main()

固定長のbit列を単純に割り当てて復号化した場合と比較すると文字列全体の長さは69%程度となっています。

input_str='DAEBCBACBBBC'
{'B': '0', 'C': '10', 'A': '110', 'E': '1110', 'D': '1111'}
encoded_str='1111110111001001101000010'
rate=0.6944444444444444

参考

Discussion

ログインするとコメントできます