Untitled

 avatar
unknown
plain_text
2 years ago
1.5 kB
2
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)