Huffman.py

class HuffmanTree(object):
    def __init__(self, left, right):
        self.left = left
        self.right = right
    
    def encode(self, symbol):
        for branch, instruction  in ((self.left, "0"), (self.right, "1")):
            if type(branch) is HuffmanTree:
                if branch.encode(symbol):
                    return instruction + branch.encode(symbol)
            elif branch == symbol:
                return instruction
    
    def decode_message(self, instructions):
        message = ""
        tree = self
        for instruction in instructions:
            tree = tree.left if instruction == "0" else tree.right
            if type(tree) is str:
                message += tree
                tree = self
        return message
    
    def encode_message(self, message):
        return "".join([self.encode(char) for char in message])
        
        
argmin = lambda x: min(x, key=x.get)

def pop_least_frequent_symbol(weights):
    symbol = argmin(weights)
    weight = weights[symbol]
    del weights[symbol]
    return symbol, weight

def buildHuffmanTree(weights):
    while len(weights) > 1:
        last_symbol, last_weight = pop_least_frequent_symbol(weights)
        second_last_symbol, second_last_weight = pop_least_frequent_symbol(weights)
        weights[HuffmanTree(last_symbol, second_last_symbol)] = last_weight + second_last_weight
    
    [(tree, _ )] = weights.items()
    return tree


weights = {
    "1": 0.25,
    "2": 0.25,
    "3": 0.2,
    "4": 0.15,
    "5": 0.15
}

tree = buildHuffmanTree(weights)
print(tree.decode_message(tree.encode_message("5123312124")))