Huffman

import argparse


class Tree:

    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

    def isLeaf(self):
        return self.left is None and self.right is None

    def __repr__(self):
        return '{} [left: {} /\ right: {}]'.format(self.value, self.left,
                                                   self.right)


def min_tree(trees):
    return min(trees, key=lambda x: x.value)


def find_codes(tree, s):
    if tree.isLeaf():
        return [tree.value, s]
    else:
        return [find_codes(tree.left, '0' + s),
                find_codes(tree.right, '1' + s)]


def flatten(x):
    if isinstance(x, list):
        return [a for i in x for a in flatten(i)]
    else:
        return [x]


if __name__ == "__main__":
    p = argparse.ArgumentParser('Create a Huffman coding for a probability \
                                 distribution.')
    p.add_argument('data', type=str, help='Comma-separated probabilities')

    args = p.parse_args()
    data = []
    with open(args.data) as f:
        for d in f.readline().split(','):
            t = Tree(float(d))
            data.append(t)
    while len(data) > 1:
        data.sort(key=lambda x: x.value)
        t1 = min_tree(data)
        data.remove(t1)
        t2 = min_tree(data)
        data.remove(t2)
        new_tree = Tree(t1.value + t2.value)
        new_tree.left = t1
        new_tree.right = t2
        data.append(new_tree)
    data = data[0]
    print flatten(find_codes(data, ''))