Back to Code List
python

Huffman Tree

A generator for huffman trees.

   main.py


def make_leaf(symbol, weight):
    return (symbol, weight)

def is_leaf(x):
    return isinstance(x, tuple) and \
            len(x) == 2 and \
            isinstance(x[0], str) and \
            isinstance(x[1], int)

def get_leaf_symbol(leaf):
    return leaf[0]

def get_leaf_freq(leaf):
    return leaf[1]

def get_left_branch(huff_tree):
    return huff_tree[0]

def get_right_branch(huff_tree):
    return huff_tree[1]

def get_symbols(huff_tree):
    if is_leaf(huff_tree):
        return [get_leaf_symbol(huff_tree)]
    else:
        return huff_tree[2]

def get_freq(huff_tree):
    if is_leaf(huff_tree):
        return get_leaf_freq(huff_tree)
    else:
        return huff_tree[3]

def make_huffman_tree(left_branch, right_branch):
    return [left_branch,
            right_branch,
            get_symbols(left_branch) + get_symbols(right_branch),
            get_freq(left_branch) + get_freq(right_branch)]

def huffman_encode(symbols, huff_tree):
    outputArr = []
    for s in symbols:
        outputArr.extend(encode_symbol(s, huff_tree))
    return outputArr

def encode_symbol(symbol, huff_tree):
    currArr = []
    currBranch = huff_tree
    # loop down until we hit a leaf
    while not is_leaf(currBranch):
        lb = get_left_branch(currBranch)
        rb = get_right_branch(currBranch)
        if not symbol in get_symbols(currBranch):
            return "symbol doesn't exist"
        elif symbol in get_symbols(lb):
            currArr.append(0)
            currBranch = lb
        elif symbol in get_symbols(rb):
            currArr.append(1)
            currBranch = rb
        else:
            return "symbol not in branch"
    return currArr

def huffman_decode(bits, huff_tree):
    outputArr = []
    if bits == []: return []
    currBranch = huff_tree
    nextBranch = None
    # loop until we've covered every bit
    while bits != []:
        nextBranch = get_branch(bits[0], currBranch)
        if is_leaf(nextBranch):
            outputArr.append(get_leaf_symbol(nextBranch))
            currBranch = huff_tree
        else:
            currBranch = nextBranch
        del bits[0]
    return outputArr

def get_branch(bit, huff_tree):
    if bit == 0:
        return get_left_branch(huff_tree)
    elif bit == 1:
        return get_right_branch(huff_tree)
    else:
        return "not a bit"

## leaves
E_1 = make_leaf('E', 1)
F_1 = make_leaf('F', 1)

G_1 = make_leaf('G', 1)
H_1 = make_leaf('H', 1)

C_1 = make_leaf('C', 1)
D_1 = make_leaf('D', 1)

B_3 = make_leaf('B', 3)
A_8 = make_leaf('A', 8)

EF_2 = make_huffman_tree(E_1, F_1)
GH_2 = make_huffman_tree(G_1, make_leaf('H', 1))
EFGH_4 = make_huffman_tree(EF_2, GH_2)

CD_2 = make_huffman_tree(C_1, D_1)
BCD_5 = make_huffman_tree(B_3, CD_2)

BCDEFGH_9 = make_huffman_tree(BCD_5, EFGH_4)

ABCDEFGH_17 = make_huffman_tree(A_8, BCDEFGH_9)

print ABCDEFGH_17