class TreeNode:
def __init__(self, weight, left = None, right = None):
self.weight = weight
self.left = left
self.right = right
self.height = 1
# 递归计算节点的高度
def get_height(self):
if not self.left and not self.right:
return 1
left_height = self.left.get_height() if self.left else 0
right_height = self.right.get_height() if self.right else 0
return 1 + max(left_height, right_height)
# 重写方法,基于节点的权重进行比较。这解决了 heapq 无法比较 Node 实例的问题
def __lt__(self, other):
return self.weight < other.weight
import heapq
class Solution:
def build_huffman_tree(self, weights):
if not weights:
return
# 初始化最小堆,创建叶子节点
heap = [TreeNode(weight) for weight in weights]
heapq.heapify(heap)
while len(heap) > 1:
# 取出两个权值最小的节点
left = heapq.heappop(heap)
right = heapq.heappop(heap)
# 如果权值相同,按照高度进行处理
if left.weight == right.weight:
left_height = left.get_height()
right_height = right.get_height()
# 确保左子树高度 <= 右子树高度
if left_height > right_height:
left, right = right, left
# 创建新的父节点,权值为左右节点权值之和
merged_node = TreeNode(left.weight + right.weight, left, right)
heapq.heappush(heap, merged_node)
# 返回堆中唯一的节点,即哈夫曼树的根节点
return heap[0]
# 栈迭代法进行中序遍历
def midorder_traversal(self, root:TreeNode):
res = []
if root is None:
return
stack = [root]
while stack:
node = stack.pop()
if node != None:
# 按照顺序压栈:右节点、中节点、左节点
if node.right:
stack.append(node.right)
stack.append(node)
# 中节点访问过,但是还没有处理,加入空节点做为标记。
stack.append(None)
if node.left:
stack.append(node.left)
# 只有遇到空节点的时候,才将下一个节点放进结果集
else:
# 按照左节点、中节点、右节点出栈
node = stack.pop()
res.append(node.weight)
return res
# 递归法进行中序遍历
def midorder_traversal_I(self, root:TreeNode):
res = []
def dfs(node):
if node is None:
return
dfs(node.left)
res.append(node.weight)
dfs(node.right)
dfs(root)
return res
# 递归法简化写法
def midorder_traversal_II(self, root:TreeNode):
if root is None:
return []
return self.midorder_traversal_II(root.left) + [root.weight] + self.midorder_traversal_II(root.right)
if __name__=='__main__':
weights = list(map(int, input().split(',')))
solution = Solution()
huffman_tree = solution.build_huffman_tree(weights)
result = solution.midorder_traversal(huffman_tree)
print(result)
知识点:哈夫曼树、堆、二叉树的中序遍历、贪心
结语:越简单的题目解法应该越多,请路过大神留下新的思路供本小白学习一下,打开思路