哈夫曼编码的实现(python)

思路

C++实在是弄不懂

代码

from typing import List, Optional

class Node:
    """
    结点
    """
    def __init__(self):
        self.name: Optional[str] = None
        self.weight: Optional[int] = None
        self.left_child: Optional[Node] = None
        self.right_child: Optional[Node] = None
        self.parent: Optional[Node] = None

class NodeCode:
    """
    结点的编码
    """
    def __init__(self, name: str, code: str):
        self.name = name
        self.code = code

    def __str__(self):
        return f"{self.name}:{self.code}"

class HaffumanTreeCode:
    """
    哈夫曼树的编码
    """
    def __init__(self):
        self.code_list: List[NodeCode] = []

    def append(self, code: NodeCode):
        self.code_list.append(code)

    def code(self, massage: str):
        """
        针对编码表编码
        :param massage: 需要编码的信息
        :return:
        """
        code = ""
        for i in massage:
            for n in self.code_list:
                if n.name == i:
                    code += n.code
                    break
            # raise ValueError(f"'{i}'不存在")
        return code

    def decode(self, code: str):
        """
        针对编码表解码
        :param code:需要解码的哈夫曼码
        :return:
        """
        out = ""
        i = 0
        while code:
            code_need_match = code[0:i+1]
            for n in self.code_list:
                if n.code == code_need_match:
                    out += n.name
                    code = code[i+1:]
                    i = 0
                    break
            i += 1
        return out

    def __str__(self):
        """
        打印
        :return:
        """
        out = ""
        for i in self.code_list:
            out += i.__str__() + "\n"
        return out.strip()

class HaffumanTree:
    """
    哈夫曼树
    """
    def __init__(self):
        self.root = Node()
        self.nodes: List[Node] = []
        self.nodes_raw: List[Node] = []
        self.codes: HaffumanTreeCode = HaffumanTreeCode()

    def inser(self, newnode: Node):
        for i in range(len(self.nodes)):
            if self.nodes[i].weight < newnode.weight:
                self.nodes.insert(i, newnode)
                return
        self.nodes.append(newnode)

    def build_tree(self, name_list: List[str], weight_list: List[int]):
        """
        创建哈夫曼树
        :param name_list: 字符列表
        :param weight_list: 权重列表
        :return:
        """
        # 队每个字符生成结点
        while len(name_list) > 0:
            max_weight = max(weight_list)
            max_index = weight_list.index(max_weight)
            max_name = name_list[max_index]
            newnode = Node()
            newnode.name = max_name
            newnode.weight = max_weight
            self.nodes.append(newnode)
            name_list.remove(max_name)
            weight_list.remove(max_weight)

        # 保存源结点
        self.nodes_raw = self.nodes[:]

        # 创建新结点并生成树
        while len(self.nodes) > 2:
            left = self.nodes.pop()
            right = self.nodes.pop()
            newnode = Node()
            newnode.weight = left.weight + right.weight
            newnode.left_child = left
            newnode.right_child = right
            left.parent = newnode
            right.parent = newnode
            # 根据权重插入
            self.inser(newnode)

        left = self.nodes.pop()
        right = self.nodes.pop()
        self.root.weight = left.weight + right.weight
        self.root.left_child = left
        self.root.right_child = right
        left.parent = self.root
        right.parent = self.root

    def make_code(self):
        """
        生成编码表
        :return:
        """
        while len(self.nodes_raw):
            node = self.nodes_raw.pop()
            name = node.name
            code = ""
            while node.parent != self.root:
                parent = node.parent
                if parent.left_child == node:
                    code = "0" + code
                elif parent.right_child == node:
                    code = "1" + code
                node = parent
            if self.root.left_child == node:
                code = "0" + code
            elif self.root.right_child == node:
                code = "1" + code
            self.codes.append(NodeCode(name, code))

    def decode(self, code: str):
        return self.codes.decode(code)

    def code(self, massage: str):
        return self.codes.code(massage)

if __name__ == '__main__':
    h = HaffumanTree()
    name_list = ["a", "b", "c", "d", "e"]
    wight_list = [5, 11, 6, 21, 19]
    h.build_tree(name_list, wight_list)
    h.make_code()
    print(f"编码表：\n{h.codes}")
    print(f'bcacce编码为：{h.code("bcacce")}')
    print(f'01001000001001解码为： {h.decode("0100100000100110")}')