Huffman Encoding

huffman_encoding.py

from collections import defaultdict
import heapq

class Node:
    def __init__(self, freq, value, parent, left=None, right=None):
        self.freq = freq
        self.value = value
        self.left = left
        self.right = right

    def __lt__(self, other):
        return self.freq < other.freq

    def __eq__(self, other):
        return self.freq == other.freq

    def __gt__(self, other):
        return self.freq > other.freq

def pre_order_traversal(root, prefix, freq_map):
    # base conditions
    if root is None:
        return
    if root.left is None and root.left is None:
        freq_map[root.value] = prefix
        return
    pre_order_traversal(root.left, prefix + "0", freq_map)
    pre_order_traversal(root.right, prefix + "1", freq_map)

def huffman_encoding(string):
    # compute frequency map
    freq = defaultdict(int)
    for each_char in string:
        freq[each_char] += 1
    heap = []
    for each_char, frequency in freq.items():
        heapq.heappush(heap, Node(frequency, each_char, None))
    root = None
    while len(heap) > 1:
        # pop the first and second minimum from heap
        first_min = heapq.heappop(heap)
        second_min = heapq.heappop(heap)
        new_node = Node(first_min.freq + second_min.freq, "*", None, first_min, second_min)
        first_min.parent = new_node
        second_min.parent = new_node
        heapq.heappush(heap, new_node)
        root = new_node
    
    encoding_map = {}
    pre_order_traversal(root, "",  encoding_map)
    print(encoding_map)
    encoded_str = ""
    for each_char in string:
        encoded_str += encoding_map[each_char]
    return root, encoded_str


def huffman_decoding(root, encoded_str):
    # iterat through encoded string
    # if 1 go right, if zero go left... do until there is leaf...
    # if we go to leaf, going back to parent would be hard as we are not maintaing parent pointer
    # added parent node for each node
    # once leaf is hit, we go back to the parent
    # maintain a visited array as well...
    # once a leaf node is hit, go back
    # once both the leafs are visited, go back
    
    # once you hit leaf node, start again from root, not from current node
    curr_node = root
    result = ""
    for each_char in encoded_str:
        if curr_node.left is None and curr_node.right is None:
            result += curr_node.value
            curr_node = root
        if each_char == '0':
            curr_node = curr_node.left
        if each_char == '1':
            curr_node = curr_node.right
    if curr_node.left is None and curr_node.right is None:
        result += curr_node.value
    return result


if __name__ == "__main__":
        root, encoded_str = huffman_encoding("lokesh sanapalli and shamhitha dhathri")
        # print(root)
        print(encoded_str)
        print(huffman_decoding(root, encoded_str))