第四章- Huffman编码

class HuffmanNode:
    def __init__(self, vm: list = None, left=None, right=None):
        """

        :type vm: list
        """
        self.val, self.weight, self.left, self.right = vm[1], vm[0], left, right

    def __lt__(self, other):
        if isinstance(other, HuffmanNode):
            return self.weight < other.weight
        if isinstance(other, int):
            return self.weight < other
        if isinstance(other, list):
            return self.weight < other[0]
        if isinstance(other, float):
            return self.weight < other

    def __gt__(self, other):
        if isinstance(other, HuffmanNode):
            return self.weight >= other.weight
        if isinstance(other, int):
            return self.weight >= other
        if isinstance(other, list):
            return self.weight >= other[0]
        if isinstance(other, float):
            return self.weight >= other


class HuffmanTree:
    def __init__(self, priority: list):
        if len(priority) == 1:
            self.tree = HuffmanNode(priority[0], None, None)
        while len(priority) > 1:
            priority = sorted(priority)

            left_w = priority.pop(0)
            right_w = priority.pop(0)

            if isinstance(left_w, HuffmanNode):
                left_node = left_w
                left_w = left_node.weight
            else:
                left_node = HuffmanNode(left_w, None, None)
                left_w = left_w[0]
            if isinstance(right_w, HuffmanNode):
                right_node = right_w
                right_w = right_node.weight
            else:
                right_node = HuffmanNode(right_w, None, None)
                right_w = right_w[0]

            root_w = left_w + right_w
            root_node = HuffmanNode([root_w, None], left_node, right_node)
            priority.append(root_node)
        self.tree = priority[0]

    def code(self, key):
        if isinstance(self.tree, HuffmanNode):
            queue = [[self.tree.left, "0"], [self.tree.right, "1"]]
        while queue:
            cur = queue.pop(0)
            if isinstance(cur, list):
                if isinstance(cur[0], HuffmanNode):
                    # print(cur[0].weight, "    ", cur[0].val)
                    if cur[0].val is None:
                        queue.append([cur[0].left, cur[1] + "0"])
                        queue.append([cur[0].right, cur[1] + "1"])
                    else:
                        if cur[0].val == key:
                            return cur[1]
        raise ValueError('key can not be code')

    def code_string(self, pre_str: str):
        pre_str = list(pre_str)
        result = ""
        while pre_str:
            cur = pre_str.pop(0)
            result += self.code(cur)
        return result


k = int(input())
table = [(None, 0)] * 26
e = input()
# print(e)
for i, val in enumerate(e):
    index = ord(val) - 65
    if table[index][0] is None:
        table[index] = [1, val]
    else:
        table[index][0] += 1

pri = []
for i, tul in enumerate(table):
    if tul[0] is not None:
        pri.append(tul)

# print(pri)

h = HuffmanTree(pri)
p = h.code_string(e)
print(p)