Fasttext - 哈夫曼树介绍

import heapq


class HuffmanNode:
    """哈夫曼树节点"""
    def __init__(self, weight, char=None):
        self.weight = weight  # 权值
        self.char = char      # 字符（叶子节点）
        self.left = None      # 左子节点
        self.right = None     # 右子节点
    
    def __lt__(self, other):
        return self.weight < other.weight


def build_huffman_tree(char_weights):
    """
    构建哈夫曼树
    :param char_weights: 字符和权值的列表 [(char, weight), ...]
    :return: 哈夫曼树的根节点
    """
    # 1. 创建叶子节点
    heap = [HuffmanNode(weight, char) for char, weight in char_weights]
    heapq.heapify(heap)
    
    # 2. 构建哈夫曼树
    while len(heap) > 1:
        # 取出权值最小的两个节点
        left = heapq.heappop(heap)
        right = heapq.heappop(heap)
        
        # 创建新节点，权值为两子节点之和
        merged = HuffmanNode(left.weight + right.weight)
        merged.left = left
        merged.right = right
        
        heapq.heappush(heap, merged)
    
    return heap[0]


def calculate_wpl(root, depth=0):
    """
    计算带权路径长度 WPL
    :param root: 哈夫曼树的根节点
    :param depth: 当前深度
    :return: WPL
    """
    if root is None:
        return 0
    
    # 叶子节点
    if root.left is None and root.right is None:
        return root.weight * depth
    
    # 递归计算左右子树
    return (calculate_wpl(root.left, depth + 1) + 
            calculate_wpl(root.right, depth + 1))


# 测试
if __name__ == "__main__":
    # 示例：字符和权值
    char_weights = [('A', 5), ('B', 2), ('C', 3), ('D', 4)]
    
    # 构建哈夫曼树
    root = build_huffman_tree(char_weights)
    
    # 计算 WPL
    wpl = calculate_wpl(root)
    print(f"带权路径长度 WPL: {wpl}")