A class for assisting Huffman encoding of a series of symbols.

HuffmanEncoder.py

class HuffmanEncoder:
    def __init__(self):
        self.table = {}

    def Tally(self, symbol):
        self.table[symbol] = 1 + self.table.get(symbol, 0)

    def Compile(self):
        if len(self.table) == 0:
            raise Exception('Huffman encoder needs to have at least one symbol.')

        # Build a binary tree that allows us to use a variable
        # number of bits to encode each symbol based on its probability.
        # Make a list of HuffmanNodes.
        # Keep it sorted in ascending order of frequency.
        tree = sorted(HuffmanNode(x[0], x[1], None, None) for x in self.table.items())

        # While there is more than one node at the top of the tree,
        # keep removing the least populated pair of items and combine
        # them into a new internal node.
        while len(tree) != 1:
            a, b, *rest = tree
            node = HuffmanNode(None, a.count + b.count, a, b)
            tree = sorted([node] + rest)

        # The single remaining node is the root node of the tree.
        return tree[0]