Алгоритм Хаффмана: сжатие данных на Python

Зачем изучать алгоритм Хаффмана?

Алгоритм Хаффмана — это фундаментальный метод сжатия данных без потерь, используемый в ZIP, MP3 и PNG. Понимание его работы поможет вам глубже разобраться в структурах данных и оптимизации хранения информации. В этой статье мы реализуем на Python полноценный инструмент сжатия на основе дерева Хаффмана.

Что такое энтропия и энтропийное кодирование?

Энтропия — это мера случайности данных. Например, слово roooar имеет низкую энтропию из-за повторяющихся символов, а 123456 — высокую. Энтропийное кодирование присваивает более короткие коды часто встречающимся символам. Для roooar можно задать коды: r → 100, o → 0, a → 101. Тогда слово займёт 12 бит вместо 48 (6 байт).

Построение дерева Хаффмана

Дерево Хаффмана строится на основе частот символов. Сначала создаётся карта частот, затем из каждого символа — листовой узел с весом, равным частоте. Далее два узла с минимальным весом объединяются во внутренний узел, вес которого равен сумме весов потомков. Процесс повторяется, пока не останется один корневой узел.

Карта частот

from collections import defaultdict

def build_frequency_map(content: bytes) -> dict[int, int]:
    """Подсчитывает количество вхождений каждого байта."""
    freq = defaultdict(int)
    for byte in content:
        freq[byte] += 1
    return freq

Узлы дерева

import abc
import heapq
from functools import total_ordering

class BaseTreeNode(abc.ABC):
    def __init__(self, weight: int, is_leaf: bool = False):
        self._weight = weight
        self._is_leaf = is_leaf

    def weight(self) -> int:
        return self._weight

    def is_leaf(self) -> bool:
        return self._is_leaf

@total_ordering
class LeafNode(BaseTreeNode):
    def __init__(self, byte_element: int, weight: int):
        super().__init__(weight, is_leaf=True)
        self._element = byte_element

    def value(self) -> int:
        return self._element

    def __lt__(self, other):
        if not isinstance(other, BaseTreeNode):
            return NotImplemented
        return self.weight() < other.weight()

class InternalNode(BaseTreeNode):
    def __init__(self, left: BaseTreeNode, right: BaseTreeNode):
        weight = left.weight() + right.weight()
        super().__init__(weight, is_leaf=False)
        self._left = left
        self._right = right

    def left(self):
        return self._left

    def right(self):
        return self._right

    def __lt__(self, other):
        if not isinstance(other, BaseTreeNode):
            return NotImplemented
        return self.weight() < other.weight()

Сборка дерева через кучу

@staticmethod
def from_frequency_map(frequency_map: dict[int, int]) -> typing.Optional['InternalNode']:
    """Строит дерево Хаффмана из карты частот."""
    if not frequency_map:
        return None
    heap = []
    for byte, freq in frequency_map.items():
        heapq.heappush(heap, LeafNode(byte_element=byte, weight=freq))
    while len(heap) > 1:
        left = heapq.heappop(heap)
        right = heapq.heappop(heap)
        heapq.heappush(heap, InternalNode(left, right))
    return heap[0]

Создание таблицы кодирования

Обходим дерево с помощью DFS: налево — добавляем 0, направо — 1. Достигнув листа, сохраняем путь как код символа.

def build_encoding_table(root: BaseTreeNode) -> dict[int, str]:
    """Строит таблицу кодирования: байт -> строка битов."""
    table = {}
    stack = [(root, [])]
    while stack:
        node, code = stack.pop()
        if node.is_leaf():
            table[node.value()] = ''.join(code)
        else:
            if node.left():
                stack.append((node.left(), code + ['0']))
            if node.right():
                stack.append((node.right(), code + ['1']))
    return table

Запись сжатых данных

Кодируем каждый байт исходного файла, объединяем биты в массив, дополняем до целого числа байт. В начало файла записываем сериализованную таблицу кодирования (с помощью pickle) и её длину.

from bitstring import BitArray
from struct import pack
import pickle

def encode_file(input_path: str, output_path: str):
    with open(input_path, 'rb') as f:
        content = f.read()
    freq_map = build_frequency_map(content)
    root = InternalNode.from_frequency_map(freq_map)
    encoding_table = build_encoding_table(root)
    # Формируем битовый поток
    bit_array = BitArray()
    # Записываем длину таблицы (2 байта) и саму таблицу
    pickled_table = pickle.dumps(encoding_table)
    bit_array.append(pack('<H', len(pickled_table)))
    bit_array.append(pickled_table)
    # Кодируем содержимое
    encoded_bits = ''.join(encoding_table[byte] for byte in content)
    encoded_content = BitArray(bin=encoded_bits)
    # Дополнение до байта
    padding = (8 - len(encoded_content) % 8) % 8
    bit_array.append(pack('<B', padding))
    bit_array.append(encoded_content)
    with open(output_path, 'wb') as f:
        bit_array.tofile(f)

Декомпрессия

Читаем таблицу, строим обратное отображение (код -> байт), затем декодируем биты. Обратите внимание на отступы при чтении.

def decode_file(input_path: str, output_path: str):
    with open(input_path, 'rb') as f:
        data = f.read()
    bit_array = BitArray(data)
    # Читаем длину таблицы
    table_len = bit_array.read('uint:16')
    # Читаем таблицу
    table_bytes = bit_array.read(table_len * 8).tobytes()
    encoding_table = pickle.loads(table_bytes)
    # Читаем отступ
    padding = bit_array.read('uint:8')
    # Читаем закодированные данные (оставшиеся биты минус отступ)
    encoded_bits = bit_array.read(bit_array.len - bit_array.pos - padding).bin
    # Строим обратную таблицу
    reverse_table = {code: byte for byte, code in encoding_table.items()}
    # Декодируем
    current_code = ''
    decoded_bytes = []
    for bit in encoded_bits:
        current_code += bit
        if current_code in reverse_table:
            decoded_bytes.append(reverse_table[current_code])
            current_code = ''
    with open(output_path, 'wb') as f:
        f.write(bytes(decoded_bytes))

Практический вывод

Алгоритм Хаффмана — отличный способ понять, как работают современные архиваторы. Попробуйте сжать текстовый файл с помощью этого кода и сравните размеры. Для углублённого изучения посмотрите реализацию на GitHub и GitLab.

Не останавливайтесь на этом — примените знания в своих проектах! Например, реализуйте сжатие для логов или конфигурационных файлов.