树结构的python实现（五：哈夫曼树）

以下为哈夫曼树的python实现：

from collections import Counter
from binary_tree import Node, BinaryTree  # binary_tree的代码见文章：https://blog.youkuaiyun.com/moyao_miao/article/details/136787981


class HuffmanNode(Node):
    """哈夫曼树节点类"""
    def __init__(self, key_value):
        """重写构造函数，输入由键值对构成"""
        super().__init__(key_value[1])
        self.key = key_value[0]  # 编码的键


class HuffmanTree(BinaryTree):
    """哈夫曼树类"""
    def __init__(self, data_dict):
        """
        构造函数，使用给定的数据字典初始化哈夫曼树。
        :param data_dict:包含数据项及其对应权重的字典。
        """
        self.node_list = [HuffmanNode(item) for item in data_dict.items()]  # 创建节点列表
        self.root = self._list_to_binarytree()  # 构建哈夫曼树
        self.code_table = {}  # 初始化编码表

    def _list_to_binarytree(self):
        """将节点列表转换为哈夫曼树：使用贪心算法，通过不断选择两个最小权重的节点合并，构建新的哈夫曼树。"""
        while len(self.node_list) > 1:
            self.node_list.sort(key=lambda x: x.value)  # 根据节点权重排序
            left_node = self.node_list.pop(0)  # 弹出最小权重的节点作为左子节点
            right_node = self.node_list.pop(0)  # 弹出次小权重的节点作为右子节点
            node = HuffmanNode((None, left_node.value + right_node.value))  # 创建新节点，其权重为两子节点之和
            node.add(True, left_node)  # 将弹出的节点作为新节点的子节点
            node.add(False, right_node)
            self.node_list.insert(0, node)  # 将新节点添加回列表
        return self.node_list[0]  # 当列表中只剩一个元素时，即为树的根节点

    def _preorder(self, node, code=''):
        """
        重写前序遍历的内部递归函数，使其遍历并记录各叶子节点的编码到编码表。
        :param node: 当前节点
        :param code: 编码
        :return: None
        """
        if node.left_node and node.right_node:
            self._preorder(node.left_node, code + '0')
            self._preorder(node.right_node, code + '1')
        else:self.code_table[node.key] = code

    def get_code_table(self):
        """前序遍历哈夫曼树获取其编码表"""
        self.preorder_traversal()
        return self.code_table


class HuffmanCode():
    def __init__(self, data_list):
        """
        初始化哈夫曼编码器。
        :param data_list:需要被编码的原始数据列表。
        属性:
        code_table (dict): 由哈夫曼树生成的编码表，键为原始数据，值为对应的哈夫曼编码。
        encode_data (str): 原始数据列表经过哈夫曼编码后生成的二进制字符串。
        """
        self.code_table = HuffmanTree(Counter(data_list)).get_code_table()  # 使用Counter统计data_list中元素的频率，然后基于这些频率生成哈夫曼树和对应的编码表
        self.encode_data = ''.join([self.code_table[data] for data in data_list])  # 根据编码表，将data_list中每个元素转换为对应的哈夫曼编码，并拼接成一个字符串

    def encode(self):
        """
        对已生成的二进制编码串进行填充，使其长度成为8的倍数，以便于存储和传输。
        :return: 经过填充处理的二进制编码字符串，最后8位用于记录填充的0的个数。
        """
        zero_num = 8 - len(self.encode_data) % 8  # 计算为了使二进制字符串长度达到8的倍数，需要补充的0的个数
        self.encode_data += '0' * zero_num + '{:08b}'.format(zero_num)  # 在二进制字符串末尾添加相应数量的0，以及一个8位二进制数表示添加的0的个数
        return self.encode_data
        # return ''.join([hex(int(self.encode_data[8*i:8*(i+1)],2)) for i in range(len(self.encode_data)//8)])#转十六进制的方法

    def decode(self):
        """
        解码方法，将存储在实例中的二进制编码字符串还原回原始数据列表。
        处理流程包括：
        1. 从二进制编码字符串尾部读取用于填充的位数信息并移除这些和对应的填充位。
        2. 使用编码表的逆映射（解码表）将二进制编码字符串还原成原始数据。

        ### 改进建议
        1. **效率优化**：当前的解码过程是通过逐字符查找解码表来实现，这在编码数据较长时可能效率不高。考虑优化数据结构以加快查找速度，例如使用更高效的查找算法或数据结构。
        :return:从二进制编码字符串解码得到的原始数据，元素之间无分隔符连续连接。
        """
        zero_num = int(self.encode_data[-8:], 2)  # 从编码字符串末尾的8位读取用于填充的位数
        self.encode_data = self.encode_data[:-8 - zero_num]  # 移除填充的和记录填充位数的数据，获得实际编码内容
        decode_table = {value: key for key, value in self.code_table.items()}  # 构造解码表，即编码表的键值对调换
        index, code, decode_data = 0, '', ''  # 初始化指针、临时代码存储器和解码结果容器
        while index < len(self.encode_data):
            code += self.encode_data[index]
            # 如果临时存储的代码在解码表中，解码对应的字符，并重置临时存储器
            if code in decode_table:
                decode_data += decode_table[code]
                code = ''
            index += 1
        return decode_data


if __name__ == "__main__":
    obj = HuffmanTree({'a': 13, 'b': 10, 'c': 23, 'd': 6, 'e': 76, 'f': 18, 'g': 3, 'h': 52, 'i': 80})
    obj.plot()
    obj = HuffmanCode('can you can a can as a can canner can a can?')
    print(obj.encode())
    print(obj.decode())

输出：

              281
            125      156
          52     73     76     80
        N    N    31    42    N    N    N    N
      13   18   19   23
    N  N  N  N  9  10  N  N
  3 6 N N


1110100101101101101111101001011101001001101110100100111010110011011101001011101000011001011001110111010010011011101001100000000000000110
can you can a can as a can canner can a can?