huffman.py

import dataclasses
from fractions import Fraction
import math
from queue import PriorityQueue

from typing import *


@dataclasses.dataclass
class InformationSourceElement:
    alphabet: str
    probability: float

InformationSource = List[InformationSourceElement]


def merge_information_source(s1: InformationSource, s2: InformationSource) -> InformationSource:
    return [
        InformationSourceElement(
            e1.alphabet + e2.alphabet,
            e1.probability * e2.probability) for e1 in s1 for e2 in s2
    ]


def information_source_expansion(s: InformationSource, d: int) -> InformationSource:
    S = [InformationSourceElement("", 1)]

    for _ in range(d):
        S = merge_information_source(S, s)

    return S


def calculate_entropy(s: InformationSource) -> float:
    total = 0
    for e in s:
        if e.probability != 0:
            p = e.probability
            total += -p * math.log2(p)
    return total


def __flatten(tree):
    l = []

    if type(tree[0]) == str:
        l += [(tree[0], "0")]
    else:
        l += [(e[0], "0" + e[1]) for e in __flatten(tree[0])]

    if type(tree[1]) == str:
        l += [(tree[1], "1")]
    else:
        l += [(e[0], "1" + e[1]) for e in __flatten(tree[1])]

    return l

def generate_huffman_coding(s: InformationSource) -> Dict[str, str]:
    l = [(e.probability, e.alphabet) for e in s]

    while len(l) != 1:
        l = list(reversed(sorted(l)))
        a, b = l[-2], l[-1]
        l = l[:-2] + [(a[0] + b[0], (a[1], b[1]))]

    return {e[0]: e[1] for e in __flatten(l[0][1])}


def calculate_huffman_coding_mean_length(s: InformationSource) -> Fraction:
    H = generate_huffman_coding(s)

    total = 0
    for e in s:
        c = H[e.alphabet]
        total += e.probability * len(c)

    return total


p = Fraction(4, 10)
m = 8

Bn = {}
Bn[1] = B1 = [
    InformationSourceElement("0", p),
    InformationSourceElement("1", 1 - p),
]

for i in range(2, m + 1):
    Bn[i] = information_source_expansion(B1, i)

for i in range(1, m + 1):
    Bi = Bn[i]
    print(f"H(B^{i}):", calculate_entropy(Bi))

for i in range(1, m + 1):
    Bi = Bn[i]
    print(f"L(B^{i}):", float(calculate_huffman_coding_mean_length(Bi)))