diff options
Diffstat (limited to 'ass3/q2/header.py')
| -rw-r--r-- | ass3/q2/header.py | 29 |
1 files changed, 22 insertions, 7 deletions
diff --git a/ass3/q2/header.py b/ass3/q2/header.py index a10b9a9..24964af 100644 --- a/ass3/q2/header.py +++ b/ass3/q2/header.py @@ -2,6 +2,7 @@ import heapq import sys +# Represents a node in the tree used to construct a Huffman code class HuffNode: def __init__(self, string, count): self.string = string @@ -19,6 +20,9 @@ class HuffNode: new_node.height += 1 + max(self.height, other.height) return new_node + # Used to determine a node's place in the heap + # Sorts based on number of occurrences first, height second. + # This allows for a shorter code word. def __lt__(self, other): if self.count == other.count: return self.height < other.height @@ -30,6 +34,7 @@ class HuffNode: def num_children(self): return (self.left is not None) + (self.right is not None) + # Apply a function to all child nodes via DFS search def dfs_apply(self, func): func(self) if self.left is not None: @@ -38,6 +43,7 @@ class HuffNode: self.right.dfs_apply(func) +# Counts the number of unique characters in a string and places them in a dictionary def count_unique(string): count = {} for char in string: @@ -48,24 +54,28 @@ def count_unique(string): return count +# Given a string, returns a dictionary mapping from a string to a binary code def huffman(string): count = count_unique(string) alphabet = "".join(count.keys()) - nodes = [HuffNode(letter, count[letter]) for letter in alphabet] - leaves = nodes[:] + nodes = [HuffNode(letter, count[letter]) for letter in alphabet] # Node for each letter + leaves = nodes[:] # Just the leaf nodes used later for extracting leaf Huffman codes heapq.heapify(nodes) + # Keep linking nodes together until there is only one left while len(nodes) > 1: - first, second = heapq.heappop(nodes), heapq.heappop(nodes) - new_node = first + second - heapq.heappush(nodes, new_node) + first, second = heapq.heappop(nodes), heapq.heappop(nodes) # Pop two off the heap + new_node = first + second # Combine the nodes into a new, combined node. (Non-destructive) + heapq.heappush(nodes, new_node) # Add this new node to the heap root = nodes[0] - huffman_calculate(root) - result = {} + huffman_calculate(root) # Calculate Huffman codes for the entire tree + result = {} # Combine results into a dictionary for leaf in leaves: result[leaf.string] = leaf.binary return result +# Calculates the Huffman codes of this node and all of its children recursively +# Non-leaf nodes are calculated as intermediate values. Only the leaves are considered important def huffman_calculate(node: HuffNode, bin_prefix=""): if node.num_children() == 0: node.binary = bin_prefix @@ -74,6 +84,7 @@ def huffman_calculate(node: HuffNode, bin_prefix=""): huffman_calculate(node.right, bin_prefix + "1") +# Calculate the Elias omega code for a given positive integer def elias(n): assert n > 0 binary = bin(n)[2:] @@ -84,6 +95,7 @@ def elias(n): return result +# Generate the binary header for a given string def header(string): huff_code = huffman(string) alphabet_size = len(huff_code.keys()) @@ -95,12 +107,14 @@ def header(string): return output +# Read a file and return the contents def read_file(filename): with open(filename, "r") as file: contents = file.read() return contents +# Write a given string to file def write_file(filename, contents): with open(filename, "w") as file: file.write(contents) @@ -111,5 +125,6 @@ def main(): string = read_file(sys.argv[1]) write_file("output_header.txt", header(string)) + if __name__ == "__main__": main() |