gistfile1.py

"Calculates the binary huffman code for the given characters and their frequencies."

from copy import copy


freqs = {'a': 20, 'b': 10, 'c': 12, 'd': 5, 'e': 15, 'z': 65 }


class Node:
	def __init__(self, ch='', freq=-1, code=None):
		self.ch = ch
		self.freq = freq
		self.code = code
		self.left = None
		self.right = None
	
	def __cmp__(self, other):
		return self.freq - other.freq
	

def get_huffman_tree(freqs):
	"Returns the Node of the huffman code tree, with each character having 0 or 1."

	nodelist = []

	# set up nodes
	for ch,f in freqs.items():
		nodelist.append(Node(ch, f))
	

	while len(nodelist) > 1:
		nodelist.sort(reverse=True)
		l = nodelist.pop()
		r = nodelist.pop()
		
		l.code = 0
		r.code = 1
		
		tmpNode = Node(freq=l.freq + r.freq)
		tmpNode.left = l
		tmpNode.right = r

		nodelist.append(tmpNode)

	if nodelist:
		return nodelist.pop()
	
	return None


def print_node(node):
	print node.ch, node.code


def prefix_traverse(tree, f, codes=[]):
	"Traverses the tree, printing characters and their code."
	if tree:
		# get the current co
		if tree.code != None:
			codes.append(tree.code)
		
		# if leaf, print ch & code
		if tree.ch:
			print tree.ch, "".join([str(c) for c in codes])

		prefix_traverse(tree.left, f, copy(codes))
		prefix_traverse(tree.right, f, codes)



if __name__ == "__main__":
	
	huff_tree = get_huffman_tree(freqs)

	prefix_traverse(huff_tree, print_node)