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.'''
    assert(sum(p.values()) == 1.0) # Ensure probabilities sum to 1

    # 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)
    p_prime[a1 + a2] = p1 + p2

    # Recurse and construct code on new distribution
    c = huffman(p_prime)
    ca1a2 = c.pop(a1 + a2)
    c[a1], c[a2] = ca1a2 + '0', 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 (i,pi): pi)
    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'}

A trivial question. How you know which is 0 and which is 1 in p when len(p) == 2? I think a sorting should be applied here.

It doesn't matter which one is which because they both has the same codeword length of 1.

While true, for practical usage we may want reproducible results. Your example execution can generate different mappings across runs or across different python versions, since the dictionaries will have different orderings.

Here's a patch that overcomes those issues by sorting dictionaries:

--- huffman.py
+++ huffman_sorted.py
@@ -8,7 +8,7 @@
 
     # Base case of only two symbols, assign 0 or 1 arbitrarily
     if(len(p) == 2):
-        return dict(zip(p.keys(), ['0', '1']))
+        return dict(list(zip(list(sorted(p.keys())), ['0', '1'])))
 
     # Create a new distribution by merging lowest prob. pair
     p_prime = p.copy()
@@ -27,7 +27,7 @@
     '''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 (i,pi): pi)
+    sorted_p = sorted(list(p.items()), key=lambda i_pi: (i_pi[1],i_pi[0]))
     return sorted_p[0][0], sorted_p[1][0]
 
 # Example execution