Huffman Coding in Python

import heapq
from collections import defaultdict, Counter

class HuffmanNode:
    def __init__(self, symbol, frequency):
        self.symbol = symbol
        self.frequency = frequency
        self.left = None
        self.right = None

    def __lt__(self, other):
        return self.frequency < other.frequency

def build_huffman_tree(data):
    frequency_counter = Counter(data)
    heap = [HuffmanNode(symbol, freq) for symbol, freq in frequency_counter.items()]
    heapq.heapify(heap)

    while len(heap) > 1:
        left_child = heapq.heappop(heap)
        right_child = heapq.heappop(heap)

        internal_node = HuffmanNode(None, left_child.frequency + right_child.frequency)
        internal_node.left = left_child
        internal_node.right = right_child

        heapq.heappush(heap, internal_node)

    return heap[0]

def build_huffman_codes(node, current_code, codes):
    if node.symbol is not None:
        codes[node.symbol] = current_code
        return

    if node.left is not None:
        build_huffman_codes(node.left, current_code + '0', codes)

    if node.right is not None:
        build_huffman_codes(node.right, current_code + '1', codes)

def huffman_encoding(data):
    if not data:
        return None, None

    root = build_huffman_tree(data)
    codes = {}
    build_huffman_codes(root, '', codes)

    encoded_data = ''.join(codes[symbol] for symbol in data)
    return encoded_data, root

def huffman_decoding(encoded_data, root):
    if not encoded_data or root is None:
        return None

    current_node = root
    decoded_data = ''

    for bit in encoded_data:
        if bit == '0':
            current_node = current_node.left
        else:
            current_node = current_node.right

        if current_node.symbol is not None:
            decoded_data += current_node.symbol
            current_node = root

    return decoded_data

# Example Usage:
data_to_encode = "hello, world!"
encoded_data, huffman_tree = huffman_encoding(data_to_encode)

print("Original data:", data_to_encode)
print("Encoded data:", encoded_data)

decoded_data = huffman_decoding(encoded_data, huffman_tree)
print("Decoded data:", decoded_data)