python huffman.py source.txt result.txt
1 import sys
2 import heapq
3 import collections
4
5 class Node(object):
6 def __init__(self,value = None,count = 1,left = None,right = None, code = ''):
7 self.value = value
8 self.count = count
9 self.left = left
10 self.right = right
11 self.code = code
12
13 def isleaf(self):
14 if self.left != None:
15 return False
16 elif self.right != None:
17 return False
18 else:
19 return True
20
21 def __repr__(self):
22 return "Node(%r,%r)"%(self.value,self.count)
23 #for sort or priority queue
24 def __lt__(self,other):
25 return self.count < other.count
26 #for operator +
27 def __add__(self,other):
28 self.code = 0
29 other.code = 1
30 # only leaf node value is need
31 return Node(self.value+other.value,self.count+other.count,self,other)
32
33 def getTreeRoot(text):
34 Counter = collections.Counter(text)
35 head = [Node(k,v) for (k,v) in Counter.items()]
36 heapq.heapify(head)
37
38 while len(head) >= 2:
39 heapq.heappush(head, heapq.heappop(head) + heapq.heappop(head))
40
41 root = head[0]
42 return root
43
44 def huffman(root, prefix = []):
45 code = {}
46 if root is None:
47 return code
48 prefix = prefix + [root.code]
49 if root.isleaf():
50 code[root.value] = prefix
51 else:
52 code.update(huffman(root.left,prefix))
53 code.update(huffman(root.right,prefix))
54 return code
55
56 def gethuffmantext(text):
57 root = getTreeRoot(text)
58 codebook = huffman(root)
59 for k,v in codebook.items():
60 newv = "".join(str(char) for char in v)
61 codebook[k] = newv
62 print (codebook)
63 print("The original text size is {}".format(len(text) * 8))
64 huffmantext = []
65 lenhuffman = 0
66 for char in text:
67 lenhuffman += len(codebook[char])
68 huffmantext.append(codebook[char])
69 print ("The huffman code text size is {}".format(lenhuffman))
70
71 return huffmantext, codebook, lenhuffman
72
73 def gettextfromhuffmancode(huffmantext,codebook):
74
75 reversecodebook = {value:key for (key,value) in codebook.items()}
76
77 text = []
78 for huffmancode in huffmantext:
79 text.append(reversecodebook[huffmancode])
80
81 return text
82
83
84 if __name__ == "__main__":
85 text = open(sys.argv[1],"rb").read()
86 newfile = sys.argv[2]
87 #huffmantext, codebook, lenhuffman = gethuffmantext("hello world")
88 huffmantext, codebook, lenhuffman = gethuffmantext(text)
89 text1 = gettextfromhuffmancode(huffmantext,codebook)
90 text1 = "".join(str(v) for v in text1)
91
92 lenoftext = len(text)*8.0
93 lenofhuffman = lenhuffman
94
95 print ("The compression ratio is %lf \n" % (1.0 * lenhuffman / lenoftext ))
96
97 fp = open(newfile,"wb")
98 fp.write(text1)
99 fp.close()
