Add huffman coding (encode and decode file), add tests for huffman coding (keon#500)

Author: Ziken

README.md

@@ -113,6 +113,8 @@ If you want to uninstall algorithms, it is as simple as:
     - [binary_gap](algorithms/bit/binary_gap.py)
 - [calculator](algorithms/calculator)
     - [math_parser](algorithms/calculator/math_parser.py)
+- [compression](algorithms/compression)
+    - [huffman_coding](algorithms/compression/huffman_coding.py)
 - [dfs](algorithms/dfs)
     - [all_factors](algorithms/dfs/all_factors.py)
     - [count_islands](algorithms/dfs/count_islands.py)

algorithms/compression/__init__.py

@@ -0,0 +1,328 @@
+"""
+Huffman coding is an efficient method of compressing data without losing information.
+This algorithm analyzes the symbols that appear in a message.
+Symbols that appear more often will be encoded as a shorter-bit string
+while symbols that aren't used as much will be encoded as longer strings.
+"""
+
+from collections import defaultdict, deque
+import heapq
+
+
+class Node:
+    def __init__(self, frequency=0, sign=None, left=None, right=None):
+        self.frequency = frequency
+        self.sign = sign
+        self.left = left
+        self.right = right
+
+    def __lt__(self, other):
+        return self.frequency < other.frequency
+
+    def __gt__(self, other):
+        return self.frequency > other.frequency
+
+    def __eq__(self, other):
+        return self.frequency == other.frequency
+
+    def __str__(self):
+        return "<ch: {0}: {1}>".format(self.sign, self.frequency)
+
+    def __repr__(self):
+        return "<ch: {0}: {1}>".format(self.sign, self.frequency)
+
+
+class HuffmanReader:
+    def __init__(self, file):
+        self.file = file
+        self.buffer = []
+        self.is_last_byte = False
+
+    def get_number_of_additional_bits_in_the_last_byte(self) -> int:
+        bin_num = self.get_bit() + self.get_bit() + self.get_bit()
+        return int(bin_num, 2)
+
+    def load_tree(self) -> Node:
+        """
+        Load tree from file
+
+        :return:
+        """
+        node_stack = deque()
+        queue_leaves = deque()
+        root = Node()
+
+        current_node = root
+        is_end_of_tree = False
+        while not is_end_of_tree:
+            current_bit = self.get_bit()
+            if current_bit == "0":
+                current_node.left = Node()
+                current_node.right = Node()
+                node_stack.append(current_node.right)  # going to left node, right push on stack
+                current_node = current_node.left
+            else:
+                queue_leaves.append(current_node)
+                if node_stack:
+                    current_node = node_stack.pop()
+                else:
+                    is_end_of_tree = True
+
+        self._fill_tree(queue_leaves)
+
+        return root
+
+    def _fill_tree(self, leaves_queue):
+        """
+        Load values to tree after reading tree
+        :param leaves_queue:
+        :return:
+        """
+        leaves_queue.reverse()
+        while leaves_queue:
+            node = leaves_queue.pop()
+            s = int(self.get_byte(), 2)
+            node.sign = s
+
+    def _load_byte(self, buff_limit=8) -> bool:
+        """
+        Load next byte is buffer is less than buff_limit
+        :param buff_limit:
+        :return: True if there is enough bits in buffer to read
+        """
+        if len(self.buffer) <= buff_limit:
+            byte = self.file.read(1)
+
+            if not byte:
+                return False
+
+            i = int.from_bytes(byte, "big")
+            self.buffer.extend(list("{0:08b}".format(i)))
+
+        return True
+
+    def get_bit(self, buff_limit=8):
+        if self._load_byte(buff_limit):
+            bit = self.buffer.pop(0)
+            return bit
+        else:
+            return -1
+
+    def get_byte(self):
+        if self._load_byte():
+            byte_list = self.buffer[:8]
+            self.buffer = self.buffer[8:]
+
+            return "".join(byte_list)
+        else:
+            return -1
+
+
+class HuffmanWriter:
+    def __init__(self, file):
+        self.file = file
+        self.buffer = ""
+        self.saved_bits = 0
+
+    def write_char(self, char):
+        self.write_int(ord(char))
+
+    def write_int(self, num):
+        bin_int = "{0:08b}".format(num)
+        self.write_bits(bin_int)
+
+    def write_bits(self, bits):
+        self.saved_bits += len(bits)
+
+        self.buffer += bits
+
+        while len(self.buffer) >= 8:
+            i = int(self.buffer[:8], 2)
+            self.file.write(bytes([i]))
+            self.buffer = self.buffer[8:]
+
+    def save_tree(self, tree):
+        """
+        Generate and save tree code to file
+        :param tree:
+        :return:
+        """
+        signs = []
+        tree_code = ""
+
+        def get_code_tree(T):
+            nonlocal tree_code
+            if T.sign is not None:
+                signs.append(T.sign)
+            if T.left:
+                tree_code += "0"
+                get_code_tree(T.left)
+            if T.right:
+                tree_code += "1"
+                get_code_tree(T.right)
+
+        get_code_tree(tree)
+        self.write_bits(tree_code + "1")  # "1" indicates that tree ended (it will be needed to load the tree)
+        for int_sign in signs:
+            self.write_int(int_sign)
+
+    def _save_information_about_additional_bits(self, additional_bits: int):
+        """
+        Overwrite first three bits in the file
+        :param additional_bits: number of bits that were appended to fill last byte
+        :return:
+        """
+        self.file.seek(0)
+        first_byte_raw = self.file.read(1)
+        self.file.seek(0)
+        first_byte = "{0:08b}".format(int.from_bytes(first_byte_raw, "big"))
+        # overwrite first three bits
+        first_byte = first_byte[3:]
+        first_byte = "{0:03b}".format(additional_bits) + first_byte
+
+        self.write_bits(first_byte)
+
+    def close(self):
+        additional_bits = 8 - len(self.buffer)
+        if additional_bits != 8:  # buffer is empty, no need to append extra "0"
+            self.write_bits("0" * additional_bits)
+            self._save_information_about_additional_bits(additional_bits)
+
+
+class TreeFinder:
+    """
+    Class to help find signs in tree
+    """
+
+    def __init__(self, tree):
+        self.root = tree
+        self.current_node = tree
+        self.found = None
+
+    def find(self, bit):
+        """
+        Find sign in tree
+        :param bit:
+        :return: True if sign is found
+        """
+        if bit == "0":
+            self.current_node = self.current_node.left
+        elif bit == "1":
+            self.current_node = self.current_node.right
+        else:
+            self._reset()
+            return True
+
+        if self.current_node.sign is not None:
+            self._reset(self.current_node.sign)
+            return True
+        else:
+            return False
+
+    def _reset(self, found=""):
+        self.found = found
+        self.current_node = self.root
+
+
+class HuffmanCoding:
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def decode_file(file_in_name, file_out_name):
+        with open(file_in_name, "rb") as file_in, open(file_out_name, "wb") as file_out:
+            reader = HuffmanReader(file_in)
+            additional_bits = reader.get_number_of_additional_bits_in_the_last_byte()
+            tree = reader.load_tree()
+
+            HuffmanCoding._decode_and_write_signs_to_file(file_out, reader, tree, additional_bits)
+
+        print("File decoded.")
+
+    @staticmethod
+    def _decode_and_write_signs_to_file(file, reader: HuffmanReader, tree: Node, additional_bits: int):
+        tree_finder = TreeFinder(tree)
+        is_end_of_file = False
+
+        while not is_end_of_file:
+            bit = reader.get_bit()
+            if bit != -1:
+                while not tree_finder.find(bit):  # read whole code
+                    bit = reader.get_bit(0)
+                file.write(bytes([tree_finder.found]))
+            else:  # There is last byte in buffer to parse
+                is_end_of_file = True
+                last_byte = reader.buffer
+                last_byte = last_byte[:-additional_bits]  # remove additional "0" used to fill byte
+                for bit in last_byte:
+                    if tree_finder.find(bit):
+                        file.write(bytes([tree_finder.found]))
+
+    @staticmethod
+    def encode_file(file_in_name, file_out_name):
+        with open(file_in_name, "rb") as file_in, open(file_out_name, mode="wb+") as file_out:
+            signs_frequency = HuffmanCoding._get_char_frequency(file_in)
+            file_in.seek(0)
+            tree = HuffmanCoding._create_tree(signs_frequency)
+            codes = HuffmanCoding._generate_codes(tree)
+
+            writer = HuffmanWriter(file_out)
+            writer.write_bits("000")  # leave space to save how many bits will be appended to fill the last byte
+            writer.save_tree(tree)
+            HuffmanCoding._encode_and_write_signs_to_file(file_in, writer, codes)
+            writer.close()
+
+        print("File encoded.")
+
+    @staticmethod
+    def _encode_and_write_signs_to_file(file, writer: HuffmanWriter, codes: dict):
+        sign = file.read(1)
+        while sign:
+            int_char = int.from_bytes(sign, "big")
+            writer.write_bits(codes[int_char])
+            sign = file.read(1)
+
+    @staticmethod
+    def _get_char_frequency(file) -> dict:
+        is_end_of_file = False
+        signs_frequency = defaultdict(lambda: 0)
+        while not is_end_of_file:
+            prev_pos = file.tell()
+            sign = file.read(1)
+            curr_pos = file.tell()
+            if prev_pos == curr_pos:
+                is_end_of_file = True
+            else:
+                signs_frequency[int.from_bytes(sign, "big")] += 1
+
+        return signs_frequency
+
+    @staticmethod
+    def _generate_codes(tree: Node) -> dict:
+        codes = dict()
+        HuffmanCoding._go_through_tree_and_create_codes(tree, "", codes)
+        return codes
+
+    @staticmethod
+    def _create_tree(signs_frequency: dict) -> Node:
+        nodes = [Node(frequency=frequency, sign=char_int) for char_int, frequency in signs_frequency.items()]
+        heapq.heapify(nodes)
+
+        while len(nodes) > 1:
+            left = heapq.heappop(nodes)
+            right = heapq.heappop(nodes)
+            new_node = Node(frequency=left.frequency + right.frequency, left=left, right=right)
+            heapq.heappush(nodes, new_node)
+
+        return nodes[0]  # root
+
+    @staticmethod
+    def _go_through_tree_and_create_codes(tree: Node, code: str, dict_codes: dict):
+        if tree.sign is not None:
+            dict_codes[tree.sign] = code
+
+        if tree.left:
+            HuffmanCoding._go_through_tree_and_create_codes(tree.left, code + "0", dict_codes)
+
+        if tree.right:
+            HuffmanCoding._go_through_tree_and_create_codes(tree.right, code + "1", dict_codes)

algorithms/compression/huffman_coding.py

@@ -0,0 +1,38 @@
+from algorithms.compression.huffman_coding import HuffmanCoding
+
+import unittest
+
+
+class TestHuffmanCoding(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.file_in_name = "huffman_coding_in.txt"
+        cls.file_out_bin_name = "huffman_coding_out.bin"
+        cls.file_out_name = "huffman_coding_out.txt"
+
+    def setUp(self):
+        import random
+        random.seed(1951)
+        with open(self.file_in_name, "wb") as file_in:
+            for i in range(10000):
+                file_in.write(bytes([random.randrange(0, 256)]))
+
+    def test_huffman_coding(self):
+        HuffmanCoding.encode_file(self.file_in_name, self.file_out_bin_name)
+        HuffmanCoding.decode_file(self.file_out_bin_name, self.file_out_name)
+
+        with open(self.file_in_name, "rb") as file_1, open(self.file_out_name, "rb") as file_2:
+            content_1 = file_1.read()
+            content_2 = file_2.read()
+
+            self.assertEqual(content_1, content_2)
+
+    def tearDown(self):
+        import os
+        os.remove(self.file_in_name)
+        os.remove(self.file_out_bin_name)
+        os.remove(self.file_out_name)
+
+
+if __name__ == "__main__":
+    unittest.main()