2020数据结构小学期（二）——哈夫曼编码算法

2、哈夫曼编码算法

输入：字符及其权值，待译码字符串，待解码字符串

功能要求：输出各字符的哈夫曼编码，输出译码字符串，输出解码字符串

源码：

class Node:
    def __init__(self, name, weight):
        self.name = name
        self.weight = weight
        self.left = None
        self.right = None
        self.father = None

    def is_left_child(self):
        return self.father.left == self


def create_prim_nodes(data_set, labels):
    if len(data_set) != len(labels):
        raise Exception('数据和标签不匹配!')
    nodes = []
    for i in range(len(labels)):
        nodes.append(Node(labels[i], data_set[i]))
    return nodes


def create_HF_tree(nodes):
    tree_nodes = nodes.copy()
    while len(tree_nodes) > 1:
        tree_nodes.sort(key=lambda node: node.weight)
        new_left = tree_nodes.pop(0)
        new_right = tree_nodes.pop(0)
        new_node = Node(None, (new_left.weight + new_right.weight))
        new_node.left = new_left
        new_node.right = new_right
        new_left.father = new_right.father = new_node
        tree_nodes.append(new_node)
    tree_nodes[0].father = None
    return tree_nodes[0]


def get_huffman_code(nodes):
    codes = {}
    for node in nodes:
        code = ''
        name = node.name
        while node.father != None:
            if node.is_left_child():
                code = '0' + code
            else:
                code = '1' + code
            node = node.father
        codes[name] = code
    return codes


def encode_huffman(codes, str_encoding):
    after_encoding = ''
    for char in str_encoding:
        flag = 0
        for key in codes.keys():
            if char == key:
                after_encoding += codes[key]
                flag = 1
                break
        if flag == 0:
            print('字符串中包含未知字符!')
            break
    return after_encoding


def decode_huffman(codes, str_decoding):
    after_decoding = ''
    temp = ''
    for char in str_decoding:
        temp += char
        if temp == str_decoding:
            print("编码存在错误！")
            break
        for key in codes.keys():
            if codes[key] == temp:
                after_decoding += key
                temp = ''
    return after_decoding


if __name__ == '__main__':
    # labels = ['a', 'b', 'c', 'd', 'e', 'f']
    # data_set = [9, 12, 6, 3, 5, 15]
    str_labels = input('请输入字符（用空格分隔）：')
    str_data = input('请输入权值（用空格分隔）：')
    str_labels = str_labels.replace(' ', '')
    list_data = str_data.split(' ')
    labels = []
    data_set = []
    for char in str_labels:
        labels.append(char)
    for list in list_data:
        data_set.append(int(list))

    # str_encoding = 'abbcefd'
    # str_decoding = '0001011001011111010'
    str_encoding = input('请输入译码字符串：')
    str_decoding = input('请输入解码字符串：')
    nodes = create_prim_nodes(data_set, labels)
    root = create_HF_tree(nodes)
    codes = get_huffman_code(nodes)
    print('各字符的哈夫曼编码：')
    for key in codes.keys():
        print(key, ': ', codes[key])

    print('译码字符串：', encode_huffman(codes, str_encoding))
    print('解码字符串：', decode_huffman(codes, str_decoding))