class TreeNode:
    def __init__(self, data, priority):
        self.data = data
        self.priority = priority
        self.code = ''
        self.left = None
        self.right = None


# 可用二叉堆，但这里用list实现最小优先队列
class MinPriorityQueue:
    def __init__(self, list):
        if list is None:
            self.list = []
        else:
            self.list = list

    def create(self, dict):
        for (d, p) in dict.items():
            self.push(TreeNode(d, p))

        return self.list

    def push(self, node):
        self.list.append(node)
        self.list.sort(key=lambda node: node.priority)
        return self.list

    def pop(self):
        return self.list.pop(0)


class HuffmanTree:
    def __init__(self):
        self.root = None
        self.dict = {}
        self.queue = []

    def createQueue(self, string):
        for char in string:
            if char not in self.dict:
                self.dict[char] = 1
            else:
                self.dict[char] += 1

        self.queue = MinPriorityQueue(None).create(self.dict)

    def createTree(self):
        n1 = self.queue.pop(0)
        n2 = self.queue.pop(0)
        self.root = TreeNode(None, n1.priority + n2.priority)
        self.root.left = n1
        self.root.right = n2
        self.queue = MinPriorityQueue(self.queue).push(self.root)


def midTraverse(root):
    stack = []
    list = []
    p = root
    while stack or p:
        while p:
            stack.append(p)
            p = p.left
        p = stack.pop()
        list.append(p)
        stack.append(p)


encodeDic = {}
decodeDic = {}


def getCode(root, x):
    if root:
        if root.data:
            encodeDic[root.data] = x
            decodeDic[x] = root.data
        getCode(root.left, x + '0')
        getCode(root.right, x + '1')


def Encode(s):
    encodeString = ''
    for char in s:
        encodeString += encodeDic[char]
    return encodeString


def Decode(s):
    decodeString = ''
    code = ''
    for char in s:
        code += char
        if code in decodeDic:
            decodeString += decodeDic[code]
            code = ''
    return decodeString


if __name__ == '__main__':
    string = 'aaabbccccd11111'
    tree = HuffmanTree()
    tree.createQueue(string)
    while tree.queue.__len__() > 1:
        tree.createTree()

    getCode(tree.root, '')
    print(encodeDic)
    print(decodeDic)
    string2 = Encode(string)
    print(string2)
    print(Decode(string2))