代码
class TreeNode(object):
def __init__(self, val, left=None, right=None):
self.val = val
self.left = left
self.right = right
def construct_tree(nodes):
if not nodes:
return None
root = TreeNode(nodes[0])
queue = [root]
index = 1
while index < len(nodes):
node = queue.pop(0)
if nodes[index] is not None:
node.left = TreeNode(nodes[index])
queue.append(node.left)
index += 1
if index < len(nodes) and nodes[index] is not None:
node.right = TreeNode(nodes[index])
queue.append(node.right)
index += 1
return root
import matplotlib.pyplot as plt
def draw_tree(root):
if not root:
return
# 获取树的高度
def get_height(node):
if not node:
return 0
return 1 + max(get_height(node.left), get_height(node.right))
height = get_height(root)
node_count = 2 ** height - 1
node_positions = {}
# 使用层序遍历确定每个节点的位置
queue = deque([(root, 0, 0)]) # (节点, 层级, 水平位置)
while queue:
node, level, x = queue.popleft()
y = height - level - 1 # 从底部开始绘制
node_positions[node] = (x, y)
if node.left:
queue.append((node.left, level + 1, x - 2 ** (height - level - 2)))
if node.right:
queue.append((node.right, level + 1, x + 2 ** (height - level - 2)))
# 绘制节点和连接线
fig, ax = plt.subplots(figsize=(10, 10))
ax.set_xlim(-node_count, node_count)
ax.set_ylim(-1, height)
ax.axis('off')
for node, (x, y) in node_positions.items():
# 绘制节点
ax.text(x, y, str(node.val), ha='center', va='center',
bbox=dict(facecolor='white', edgecolor='black', boxstyle='circle'))
# 绘制连接线
if node.left and node.left in node_positions:
x1, y1 = node_positions[node.left]
ax.plot([x, x1], [y, y1], 'black')
if node.right and node.right in node_positions:
x1, y1 = node_positions[node.right]
ax.plot([x, x1], [y, y1], 'black')
plt.show()
# 按装订区域中的绿色按钮以运行脚本。
if __name__ == '__main__':
print_hi('PyCharm')
solution = Solution()
nodes = [3,9,20,None,None,15,7,8,6,8,5,4,6,1,2]
root = construct_tree(nodes)
draw_tree(root)
效果
