OD C卷 - 建立Huffman树

生成Huffman树 (100)

• 给定一个数值数组weights，每个值代表二叉树叶子节点的权值（>=1）;
• 根据权值数组，生成哈夫曼树，并中序遍历输出；
• 左节点权值<= 右节点权值，根节点权值为左右节点权值之和；
• 左右节点权值相同时，左子树高度<=右子树;
  
  输入描述：
  第一行输入数组长度n;
  第二行输入权值，以空格分隔；
  输出描述：
  哈夫曼树的中序遍历序列，以空格分隔；

示例1
输入：
5
5 15 40 30 10
输出：
40 100 30 60 15 30 5 15 10
思路: 建立哈夫曼树的过程

• 以每个权值创建一个节点（只有一个树根的二叉树），放入nodes_list，并按照权值从小到大排序；
• 取出权值最小的两个节点，进行合并，根节点权值为左右子节点的权值之和
  • 设置根节点的左子树、右子树
  • 当左右子树（根节点）权值相同时，对比两者的树高，高的子树作为右子树；
  • 统计子树的高度 max(h_left, h_right)，h_left为每次取左子树；
• 将合并后的root根节点放入nodes_list，继续以上步骤，直到nodes_list中只有一个根节点；

# 中序遍历结果
result = []
# 输入n
n = int(input().strip())
weights = [int(x) for x in input().strip().split()]
 

class TreeNode:
    def __init__(self, left, right, weight, height):
        self.left = left
        self.right = right
        self.weight = weight
        self.height = height
 

# 中序遍历
def dfs(node):
	# 中序遍历左子树
    if node.left is not None:
        dfs(node.left)
    # 中间处理根
    result.append(node.weight)
    # 中序遍历右子树
    if node.right is not None:
        dfs(node.right)
        
    
# 建立哈夫曼树
nodes_list = []
for i in range(n):
    nodes_list.append(TreeNode(None, None, weights[i], 1))

# 比较树高
def comp(o1, o2):
    if(o1.weight == o2.weight): # 权值相等
        if o1.height > o2.height:
            return 1
        else:
            return -1
    
    if o1.weight > o2.weight:
        return 1
    else:
        return -1


while True:
    if len(nodes_list) == 1:
        break
    else :
        nodes_list = sorted(nodes_list, key=lambda i:i.weight)
        node1 = nodes_list.pop(0)
        node2 = nodes_list.pop(0)
        root = TreeNode(None, None, node1.weight + node2.weight, max(node1.height, node2.height) + 1)
        if comp(node1, node2) == 1:
        	root.right = node1
        	root.left = node2
        else:
        	root.right = node2
        	root.left = node1
        nodes_list.append(root)
    
 
 
dfs(nodes_list[0]) # 传入树根，中序遍历
# 输出字符串
output_str = ""
for i in range(len(result)):
    output_str += str(result[i])
    if(i!=len(result)-1):
        output_str+=" "
print(output_str)