Python/compression/huffman.py

83 lines
2.3 KiB
Python
Raw Normal View History

2019-05-13 05:15:27 +00:00
import sys
class Letter:
def __init__(self, letter, freq):
self.letter = letter
self.freq = freq
self.bitstring = ""
def __repr__(self):
return f'{self.letter}:{self.freq}'
class TreeNode:
def __init__(self, freq, left, right):
self.freq = freq
self.left = left
self.right = right
def parse_file(file_path):
2019-05-16 11:20:42 +00:00
"""
2019-05-13 05:15:27 +00:00
Read the file and build a dict of all letters and their
frequences, then convert the dict into a list of Letters.
"""
chars = {}
with open(file_path) as f:
while True:
c = f.read(1)
if not c:
break
chars[c] = chars[c] + 1 if c in chars.keys() else 1
2019-05-16 11:20:42 +00:00
return sorted([Letter(c, f) for c, f in chars.items()], key=lambda l: l.freq)
2019-05-13 05:15:27 +00:00
def build_tree(letters):
2019-05-16 11:20:42 +00:00
"""
2019-05-13 05:15:27 +00:00
Run through the list of Letters and build the min heap
for the Huffman Tree.
"""
while len(letters) > 1:
left = letters.pop(0)
right = letters.pop(0)
total_freq = left.freq + right.freq
node = TreeNode(total_freq, left, right)
letters.append(node)
letters.sort(key=lambda l: l.freq)
return letters[0]
def traverse_tree(root, bitstring):
2019-05-16 11:20:42 +00:00
"""
2019-05-13 05:15:27 +00:00
Recursively traverse the Huffman Tree to set each
Letter's bitstring, and return the list of Letters
"""
if type(root) is Letter:
root.bitstring = bitstring
return [root]
letters = []
letters += traverse_tree(root.left, bitstring + "0")
letters += traverse_tree(root.right, bitstring + "1")
return letters
def huffman(file_path):
2019-05-16 11:20:42 +00:00
"""
2019-05-13 05:15:27 +00:00
Parse the file, build the tree, then run through the file
2019-05-16 11:20:42 +00:00
again, using the list of Letters to find and print out the
2019-05-13 05:15:27 +00:00
bitstring for each letter.
"""
letters_list = parse_file(file_path)
root = build_tree(letters_list)
letters = traverse_tree(root, "")
print(f'Huffman Coding of {file_path}: ')
with open(file_path) as f:
while True:
c = f.read(1)
if not c:
break
le = list(filter(lambda l: l.letter == c, letters))[0]
print(le.bitstring, end=" ")
print()
if __name__ == "__main__":
# pass the file path to the huffman function
huffman(sys.argv[1])