huffman.py

#!/usr/bin/env python

from collections import defaultdict


def calc_probabilities(s="the cat in the hat"):
    """Calculate the probabilities (read: frequencies) of letters in
    string `s`.

    """
    counts = defaultdict(lambda: 0)
    for i in s:
        counts[i] += 1
    return sort_nodes(counts.items())


def sort_nodes(ns):
    """Sort nodes by their probabilities."""
    return sorted(ns, key=lambda x: x[1])


def build_node(ns):
    """Merge two nodes to build a new node."""
    a = ns[0]
    b = ns[1]
    return (sort_nodes((a, b)), a[1] + b[1])


def process_nodes(ns):
    """Create Huffman tree iteratively in the following steps:

    1. Merge two nodes with highest priority (lowest probability)
    2. Draw the rest of the fucking owl

    http://29.media.tumblr.com/tumblr_l7iwzq98rU1qa1c9eo1_500.jpg

    """
    # Iterative version of `recur_nodes`
    while len(ns) > 1:
        new_node = build_node(ns)
        ns = ns[2:]
        ns.insert(0, new_node)
        ns = sort_nodes(ns)
    return ns


def recur_nodes(ns):
    """Create Huffman tree recursively by creating Huffman tree
    recursively by creating Huffman tree recursiv

    """
    # Recursive version of `process_nodes`
    if len(ns) == 1:
        return ns

    new_node = build_node(ns)
    ns = ns[2:]
    ns.insert(0, new_node)
    ns = sort_nodes(ns)
    return recur_nodes(ns)


def main(recur=False, s=None):
    """Have you ever wanted to learn to draw owls? This function will
    draw owls for you because you can't.

    Not pictured: Generating binary codes for the nodes. Just make
    something up, no one will notice the difference LOL

    """
    nodes = calc_probabilities(s)
    nodes = sort_nodes(nodes)
    func = recur_nodes if recur else process_nodes
    return func(nodes)