Python/compression/huffman.py

89 lines
2.3 KiB
Python
Raw Normal View History

2019-05-13 05:15:27 +00:00
import sys
2019-10-05 05:14:13 +00:00
2019-05-13 05:15:27 +00:00
class Letter:
def __init__(self, letter, freq):
self.letter = letter
self.freq = freq
self.bitstring = {}
2019-05-13 05:15:27 +00:00
def __repr__(self):
2019-10-05 05:14:13 +00:00
return f"{self.letter}:{self.freq}"
2019-05-13 05:15:27 +00:00
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
frequencies, then convert the dict into a list of Letters.
2019-05-13 05:15:27 +00:00
"""
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
return sorted((Letter(c, f) for c, f in chars.items()), key=lambda l: l.freq)
2019-05-13 05:15:27 +00:00
2019-10-05 05:14:13 +00:00
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]
2019-10-05 05:14:13 +00:00
2019-05-13 05:15:27 +00:00
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 dictionary, and return the list of Letters
2019-05-13 05:15:27 +00:00
"""
if type(root) is Letter:
root.bitstring[root.letter] = bitstring
2019-05-13 05:15:27 +00:00
return [root]
letters = []
letters += traverse_tree(root.left, bitstring + "0")
letters += traverse_tree(root.right, bitstring + "1")
return letters
2019-10-05 05:14:13 +00:00
2019-05-13 05:15:27 +00:00
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
again, using the letters dictionary 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 = {
k: v for letter in traverse_tree(root, "") for k, v in letter.bitstring.items()
}
print(f"Huffman Coding of {file_path}: ")
2019-05-13 05:15:27 +00:00
with open(file_path) as f:
while True:
c = f.read(1)
if not c:
break
print(letters[c], end=" ")
2019-05-13 05:15:27 +00:00
print()
2019-10-05 05:14:13 +00:00
2019-05-13 05:15:27 +00:00
if __name__ == "__main__":
# pass the file path to the huffman function
huffman(sys.argv[1])