HuffmanCode

HuffmanCoding.py

import sys

# space replace to '_'
data={'_':18.59,'A':6.42,'B':1.27,'C':2.18,'D':3.17,'E':10.31
		,'F':2.08,'G':1.52,'H':4.67,'I':5.75,'J':0.08,'K':0.49
		,'L':3.21,'M':1.98,'N':5.72,'O':6.32,'P':1.52,'Q':0.08
		,'R':4.84,'S':5.14,'T':7.96,'U':2.28,'V':0.83,'W':1.75
		,'X':0.13,'Y':1.64,'Z':0.05,'*':0.02}


class node:
	def __init__(self,alphabet,left=None,right=None):
		self.alphabet=alphabet
		self.left=left
		self.right=right


class HuffmanCode:

	def __init__(self,data):
		self.data=data
		# sort data
		self._sorted_data=sorted(data.items(),key=lambda x:x[1],reverse=True)
		# converts data to node
		self.leafs=[]
		for i in range(len(self._sorted_data)):
			self.leafs.append(node(self._sorted_data[i][0]))
			self.leafs[i].prob=self._sorted_data[i][1]
		self.root=list(self.leafs)
		# create huffman-tree
		while True:
			select=self.root[-2:]
			self.root.remove(select[0])
			self.root.remove(select[1])
			def add_2_seq(seq,param_name):
				p1=getattr(seq[0],param_name)
				p2=getattr(seq[1],param_name)
				return p1+p2
			n=node('+',select[0],select[1])
			n.prob=add_2_seq(select,'prob');
			self.root.append(n)
			self.root.sort(key=lambda x:x.prob,reverse=True)
			# if created root then break while
			if len(self.root) is 1:
				break
		self.root=self.root[0]
		self.__node_data_extension()


	# attach rank,parent and code to node
	def __node_data_extension(self):
		def _node_data_extension(node,parent=None,code=None,offset=0):
			if node is None:
				return
			node.rank=offset
			node.parent=parent
			node.code=code
			_node_data_extension(node.left,node,0,offset+1)
			_node_data_extension(node.right,node,1,offset+1)
		_node_data_extension(self.root)




	# search max-rank and return
	def max_rank(self):
		ret=[0]
		def _max_rank(next):
			if next is None:
				return
			if ret[0]<next.rank:
				ret[0]=next.rank
			_max_rank(next.left)
			_max_rank(next.right)
		_max_rank(self.root)
		return ret[0]

	# print tree
	def print_tree(self):
		white_space=[False]
		rank=[self.max_rank()]
		def _print_tree(root,is_right=False):
			if root is None:
				if is_right:
					print('')
				white_space[0]=True
				return
			_print_tree(root.left)
			_print_tree(root.right,True)
			if white_space[0]:
				for i in range(root.rank,rank[0]):
					sys.stdout.write('       ')
				white_space[0]=False
			sys.stdout.write(" - {0}:{1:>2}".format(root.alphabet,root.rank))
		print('Code Tree : ')
		_print_tree(self.root)
		print()


	#print code
	def print_code(self):
		def _print_code(leaf):
			if leaf is None:
				return
			_print_code(leaf.parent)
			if leaf.code is not None:
				sys.stdout.write(str(leaf.code))

		print("Code :")
		for i in range(len(self.leafs)):
			sys.stdout.write(self.leafs[i].alphabet+' : ')
			_print_code(self.leafs[i])
			print('')


huf=HuffmanCode(data)
huf.print_tree()
huf.print_code()