dsa.huffman

  1""" Module to access functions for Huffman Compression. """
  2from dsa.tree import Tree, TreeNode
  3from dsa.heap import PriorityQueue
  4
  5def character_frequency(s: str) -> dict:
  6    """ 
  7    Takes a string a returns a dictionary of character frequencies.
  8
  9    Args:
 10        s (str): The string to analyze.
 11
 12    Returns:
 13        Dictionary containing character frequency.
 14    """
 15    d = {}
 16    for c in s:
 17        if c in d:
 18            d[c] += 1
 19        else:
 20            d[c] = 1
 21    return d
 22
 23def build_frequency_table(s: str) -> PriorityQueue:
 24    """ 
 25    Accepts a string to encode and returns a heap of the characters.
 26
 27    Args:
 28        s (str): The string to encode.
 29
 30    Returns:
 31        A priority queue of the characters based on frequencies.
 32    """
 33    frequency_dictionary = character_frequency(s)
 34    
 35    pq = PriorityQueue()
 36
 37    for char, count in frequency_dictionary.items():
 38        pq.push(count, TreeNode(char))#, None, None))
 39
 40    return pq
 41
 42def build_huffman_tree(pq: PriorityQueue) -> Tree:
 43    """ 
 44    Accepts a priority queue and returns a Huffman Tree.
 45
 46    Args:
 47        pq (PriorityQueue): A PriorityQueue containing TreeNodes of characters based on frequencies. 
 48
 49    Returns:
 50        A Huffman Tree.
 51    """
 52    while len(pq) > 1:
 53        priority1, node1 = pq.pop_pair()
 54        priority2, node2 = pq.pop_pair()
 55        node = TreeNode(node1.value + node2.value, node1, node2)
 56        pq.push(priority1 + priority2, node)
 57
 58    return Tree(pq.pop())
 59
 60def build_huffman_dictionary(node: TreeNode, bit_string: str="") -> dict:
 61    """
 62    Given a TreeNode, build a Huffman Dictionary.
 63
 64    Args:
 65        node (TreeNode): The Huffman Node.
 66        bit_string (str): The bit string.
 67
 68    Returns:
 69        A Huffman Dictionary.
 70    """
 71    d = {}
 72    if node.left is None and node.right is None:
 73        return {node.value: bit_string}
 74
 75    d.update(build_huffman_dictionary(node.left, bit_string + '0'))
 76    d.update(build_huffman_dictionary(node.right, bit_string + '1'))
 77
 78    return d
 79
 80def huffman_encode(st: str, hd: dict) -> str:               
 81    """
 82    Encode the string using the Huffman Dictionary.
 83
 84    Args:
 85        st (str): The string to encode.
 86        hd (dict): The Huffman Dictionary.
 87
 88    Returns:
 89        The encoded string.
 90    """
 91    s = ""
 92    for c in st:
 93        s += hd[c]
 94    return s
 95
 96def huffman_decode(encoded_data: str, tree: Tree) -> str:
 97    """
 98    Decode the encoded data using the Huffman Tree.
 99    
100    Args:
101        encoded_data (str): The encoded data.
102        tree (Tree): The Huffman Tree.
103
104    Returns:
105        The decoded data.
106    """
107    node = tree.root
108    s = ""
109    for bit in encoded_data:
110        if int(bit) == 0:
111            node = node.left
112        else:
113            node = node.right
114
115        if node.left is None and node.right is None: 
116            s += node.value
117            node = tree.root
118    return s
119
120def bitstring_to_bytes(s: str) -> bytes:
121    """
122    Convert a bitstring to bytes.
123
124    Args:
125        s (str): The bitstring.
126
127    Returns:
128        Bitstring converted to bytes.
129    """
130    return bytes(int(s[i : i + 8], 2) for i in range(0, len(s), 8))
131
132def bytes_to_bitstring(ba: bytes, bitlength: int=8) -> str:
133    """
134    Convert bytes to bitstring.
135
136    Args:
137        ba (bytes): The bytes to convert.
138        bitlength (int): The bit length.
139    
140    Returns:
141        The bytes converted to bitstring.
142    """
143    if not ba:
144        return ""
145    
146    s = ""
147    for b in ba[:-1]:
148        byte = f"{b:08b}"
149        s += byte
150    
151    byte = f"{ba[-1]:b}".zfill(bitlength) 
152    s += byte
153
154    return s