Example implementation of Huffman coding in Python

huffman.py

# Example Huffman coding implementation
# Distributions are represented as dictionaries of { 'symbol': probability }
# Codes are dictionaries too: { 'symbol': 'codeword' }

def huffman(p):
    '''Return a Huffman code for an ensemble with distribution p.'''
    # Base case of only two symbols, assign 0 or 1 arbitrarily
    if len(p) == 2:
        return dict(zip(p.keys(), ['0', '1']))

    # Create a new distribution by merging lowest prob. pair
    p_prime = p.copy()
    a1, a2 = lowest_prob_pair(p)
    p1, p2 = p_prime.pop(a1), p_prime.pop(a2)
    newkey = (a1, a2)
    p_prime[newkey] = p1 + p2

    # Recurse and construct code on new distribution
    c = huffman(p_prime)
    ca1a2 = c.pop(newkey)
    c[a1] = ca1a2 + '0'
    c[a2] = ca1a2 + '1'

    return c

def lowest_prob_pair(p):
    '''Return pair of symbols from distribution p with lowest probabilities.'''
    assert len(p) >= 2 # Ensure there are at least 2 symbols in the dist.

    sorted_p = sorted(p.items(), key=lambda kv: kv[1])
    return sorted_p[0][0], sorted_p[1][0]

# Example execution
ex1 = { 'a': 0.5, 'b': 0.25, 'c': 0.25 }
huffman(ex1) # => {'a': '0', 'c': '10', 'b': '11'}
  
ex2 = { 'a': 0.25, 'b': 0.25, 'c': 0.2, 'd': 0.15, 'e': 0.15 }
huffman(ex2)  # => {'a': '01', 'c': '00', 'b': '10', 'e': '110', 'd': '111'}