import heapq
from collections import Counter

class HuffmanNode:
    """霍夫曼树的节点"""
    def __init__(self, char, freq):
        self.char = char  # 字符，叶子节点才有效
        self.freq = freq  # 频率
        self.left = None
        self.right = None

    # 为了让堆能够比较节点大小，我们需要定义比较方法
    def __lt__(self, other):
        return self.freq < other.freq

def build_huffman_tree(text):
    """根据文本构建霍夫曼树"""
    # 1. 统计字符频率
    frequency = Counter(text)
    
    # 2. 为每个字符创建叶子节点，并放入最小堆（优先队列）
    heap = [HuffmanNode(char, freq) for char, freq in frequency.items()]
    heapq.heapify(heap)
    
    # 特殊情况：如果文本只有一种字符
    if len(heap) == 1:
        node = heapq.heappop(heap)
        # 创建一个父节点，将这个单节点作为左子树，右子树为空（或也设为相同）
        root = HuffmanNode(None, node.freq)
        root.left = node
        return root
    
    # 3. 贪心构建霍夫曼树
    while len(heap) > 1:
        # 弹出两个频率最小的节点
        left = heapq.heappop(heap)
        right = heapq.heappop(heap)
        
        # 创建一个新的内部节点，频率为两者之和
        merged = HuffmanNode(None, left.freq + right.freq)
        merged.left = left
        merged.right = right
        
        # 将新节点放回堆中
        heapq.heappush(heap, merged)
    
    # 堆中最后一个节点就是根节点
    return heap[0]

def generate_codes(node, prefix="", codebook={}):
    """递归遍历霍夫曼树，生成编码表"""
    if node is None:
        return
    
    # 如果是叶子节点，则记录其编码
    if node.char is not None:
        # 处理只有一个字符的特殊情况，确保至少有一个'0'
        codebook[node.char] = prefix if prefix else "0"
        return
    
    # 向左走加 '0'，向右走加 '1'
    generate_codes(node.left, prefix + "0", codebook)
    generate_codes(node.right, prefix + "1", codebook)
    
    return codebook

def huffman_encoding(text):
    """对文本进行霍夫曼编码"""
    if not text:
        return "", {}, None
    
    # 构建树
    root = build_huffman_tree(text)
    
    # 生成编码表
    codebook = generate_codes(root)
    
    # 编码文本
    encoded_text = ''.join(codebook[char] for char in text)
    
    return encoded_text, codebook, root

def huffman_decoding(encoded_text, root):
    """对编码后的二进制字符串进行解码"""
    if not encoded_text or not root:
        return ""
    
    decoded_chars = []
    current = root
    
    # 逐位遍历编码文本
    for bit in encoded_text:
        if bit == '0':
            current = current.left
        else:
            current = current.right
        
        # 如果到达叶子节点
        if current.char is not None:
            decoded_chars.append(current.char)
            current = root  # 重置到根节点
    
    return ''.join(decoded_chars)

# 测试
if __name__ == "__main__":
    sample_text = "Huffman coding is a greedy algorithm"
    print(f"原始文本: {sample_text}")
    print(f"原始长度 (字符数): {len(sample_text)}")
    
    # 编码
    encoded, codes, tree = huffman_encoding(sample_text)
    print(f"\n编码表:")
    for char, code in sorted(codes.items()):
        print(f"  '{char}': {code}")
    
    print(f"\n编码后的二进制串:\n{encoded}")
    print(f"编码长度 (比特数): {len(encoded)}")
    
    # 计算压缩率 (使用原始文本的ASCII编码，每个字符8比特作为基准)
    original_bits = len(sample_text) * 8
    compressed_bits = len(encoded)
    compression_ratio = (1 - compressed_bits / original_bits) * 100
    print(f"\n压缩率: {compression_ratio:.2f}%")
    
    # 解码
    decoded = huffman_decoding(encoded, tree)
    print(f"\n解码后的文本: {decoded}")
    print(f"解码是否成功: {decoded == sample_text}")