Huffman Coding

# Package: io.thecodeforge.python.dsa

import heapq
from collections import Counter

class HuffmanNode:
    def __init__(self, char, freq):
        self.char  = char
        self.freq  = freq
        self.left  = None
        self.right = None

    # Min-heap ordering by frequency
    def __lt__(self, other):
        return self.freq < other.freq

def build_huffman_tree(text):
    freq = Counter(text)

    # Min-heap of leaf nodes
    heap = [HuffmanNode(char, f) for char, f in freq.items()]
    heapq.heapify(heap)

    while len(heap) > 1:
        # Take two nodes with lowest frequency
        left  = heapq.heappop(heap)
        right = heapq.heappop(heap)

        # Merge into a parent node
        parent = HuffmanNode(None, left.freq + right.freq)
        parent.left  = left
        parent.right = right

        heapq.heappush(heap, parent)

    return heap[0]  # root of Huffman tree

text = 'aabbbccddddeeeeee'
root = build_huffman_tree(text)
print(f'Root frequency: {root.freq}')  # 17 (total characters)

def get_codes(node, prefix='', codes={}):
    if node.char is not None:  # leaf
        codes[node.char] = prefix or '0'  # handle single-char case
    else:
        get_codes(node.left,  prefix + '0', codes)
        get_codes(node.right, prefix + '1', codes)
    return codes

def huffman_encode(text):
    root  = build_huffman_tree(text)
    codes = get_codes(root, '', {})
    encoded = ''.join(codes[c] for c in text)
    return encoded, codes, root

text = 'aabbbccddddeeeeee'
encoded, codes, root = huffman_encode(text)

print('Codes:')
for char in sorted(codes):
    freq = text.count(char)
    print(f'  {repr(char)}: {codes[char]:5s}  (freq={freq}, bits={freq*len(codes[char])})')

original_bits = len(text) * 8
compressed_bits = len(encoded)
print(f'\nOriginal:   {original_bits} bits')
print(f'Compressed: {compressed_bits} bits')
print(f'Ratio:      {compressed_bits/original_bits:.1%}')