哈夫曼树可视化
import matplotlib.pyplot as plt
class Tree:
def __init__(self, weight=None, left=None, right=None):
self.weight = weight
self.left = left
self.right = right
self.is_red = False
def get_height(self): # 返回树高度，未优化算法应该比较慢
layers = [self]
layer_count = 0
while layers:
layer_count += 1
new_list = []
for node in layers:
if node.left:
new_list.append(node.left)
if node.right:
new_list.append(node.right)
layers = new_list
return layer_count
def visualize(self, axis='off'):
''' 可视化树结构
基本算法： 将树状结构映射到二维矩阵中，
如果节点左右下方有节点则把该节点加入到矩阵中的坐标中，
如节点（x,y）左下方有节点则把节点放入(x+offset,y+1)
offset为x坐标偏移量，
offset = 2**(len(matrix)-y-2)
'''
FONT_SIZE = 16
NODE_SIZE_SCALE = 1.5
figure, axes = plt.subplots(figsize=(8, 6), dpi=80)
height = self.get_height()
width_ = 2 ** (height - 1)
width = 2 * width_ + 1
matrix = [[[] for x in range(width)] for y in range(height)]
matrix[0][width_] = self # put head in the middle position
for y in range(len(matrix)):
for x in range(len(matrix[y])):
node = matrix[y][x]
if node:
x1, y1 = (1 / width) * (x + 0.5), 1 - (1 / height) * y - 0.2
axes.text(x1, y1, str(node.weight), color='white', fontsize=FONT_SIZE)
offset = 2 ** (len(matrix) - y - 2)
if node.left:
matrix[y + 1][x - offset] = node.left
x2, y2 = (1 / width) * (x - offset + 0.5), 1 - (1 / height) * (y + 1) - 0.2
line = mlines.Line2D([x1, x2], [y1, y2], zorder=-1)
axes.add_line(line)
if node.right:
matrix[y + 1][x + offset] = node.right
x2, y2 = (1 / width) * (x + offset + 0.5), 1 - (1 / height) * (y + 1) - 0.2
line = mlines.Line2D([x1, x2], [y1, y2], zorder=-1)
axes.add_line(line)
cc = plt.Circle(((1 / width) * (x + 0.5), 1 - (1 / height) * y - 0.2),
1 / width / 2 * NODE_SIZE_SCALE,
color=('r' if node.is_red else 'black'))
axes.set_aspect(1)
axes.add_artist(cc, )
plt.axis(axis)
plt.show()
class HuffmanTree:
"""哈夫曼树(huffmanTree)
构建方法：
将几点根据权重大小放到最小堆中，先从做小堆中取出两个最小元素，构成一棵树，根节点权重为两子树的和。然后将根节点插入堆中，依次上边操作，
共计操作数据长度-1次，此时堆中只有一个元素，就是哈夫曼树的根节点。WPL等于根节点的权重值。
"""
def __init__(self, element: list):
self._element = element
self._size = len(element) + 1
self._used = 0
self._case = [None for _ in range(self._size)]
def _is_full(self):
return self._used > self._size
def insert(self, root: Tree):
"""往堆中插入元素"""
if self._is_full():
return 'is full'
self._used += 1
add_node = self._used
father = self._used // 2
while father:
if root.weight > self._case[father].weight:
break
self._case[add_node] = self._case[father]
add_node = father
father //= 2
self._case[add_node] = root
def delete(self):
"""从堆中删除元素"""
if self._used == 0:
return 'is empty'
root = 1
res = self._case[root]
# node = self._used
while True:
left_son = root * 2
right_son = left_son + 1
if left_son < self._used and self._case[left_son].weight <= self._case[right_son].weight and self._case[
left_son].weight < \
self._case[self._used].weight:
self._case[root] = self._case[left_son]
root = left_son
elif right_son < self._used and self._case[right_son].weight <= self._case[left_son].weight and self._case[
right_son].weight < \
self._case[self._used].weight:
self._case[root] = self._case[right_son]
root = right_son
else:
break
self._case[root] = self._case[self._used]
self._case[self._used] = None
self._used -= 1
return res
def create(self):
"""构建哈夫曼树"""
for i in self._element:
tree = Tree(i)
self.insert(tree)
for _ in range(self._size - 2):
left = self.delete()
right = self.delete()
tree = Tree(left.weight + right.weight, left, right)
self.insert(tree)
tree = self.delete()
return tree
if __name__ == '__main__':
# heap = Heap()
a = [3, 1, 4, 5, 6]
huff = HuffmanTree(a)
r = huff.create()
# t = Tree()
r.visualize() # 可视化树结构
效果展示
