Untitled
unknown
plain_text
2 years ago
1.5 kB
5
Indexable
import matplotlib.pyplot as plt class Node: def __init__(self, val): self.val = val self.children = [] self.num_children = 0 def bfs(root): queue = [(root, 0)] while queue: node, level = queue.pop(0) node.num_children = len(node.children) + sum(child.num_children for child in node.children) for child in node.children: queue.append((child, level+1)) def dfs(node, pos, x, y, ax): if not node: return ax.text(x, y, str(node.val), fontsize=12, ha='center', va='center', bbox=dict(facecolor='white', edgecolor='black', boxstyle='circle')) for i, child in enumerate(node.children): x_child = x + (i - node.num_children / 2 + child.num_children / 2 + 0.5) * 2 ** (-y) y_child = y + 1 ax.plot([x, x_child], [y, y_child], color='blue') dfs(child, (x, y), x_child, y_child, ax) def visualize_tree(root): bfs(root) fig, ax = plt.subplots(figsize=(12, 8)) dfs(root, (0, 1), 0, 1, ax) ax.axis('off') plt.show() # przykładowe drzewo z trzema dziećmi root = Node(1) root.children.append(Node(2)) root.children.append(Node(3)) root.children.append(Node(4)) root.children[0].children.append(Node(5)) root.children[0].children.append(Node(6)) root.children[1].children.append(Node(7)) root.children[1].children.append(Node(8)) root.children[2].children.append(Node(9)) root.children[2].children.append(Node(10)) root.children[2].children.append(Node(11)) # wizualizacja drzewa visualize_tree(root)
Editor is loading...