graph_builder.py

import networkx as nx
import matplotlib.pyplot as plt
edges = [
    ("S0", "S1", "a1"),
    ("S0", "S2", "a2"),
    ("S1", "S3", "a3"),
    ("S1", "S4", "a4"),
    ("S2", "S5", "a5"),
    ("S4", "S6", "a6"),
]

nodes = {
    "S0": {"w": 4, "n": 6},
    "S1": {"w": 3, "n": 4},
    "S2": {"w": 1, "n": 2},
    "S3": {"w": 0, "n": 1},
    "S4": {"w": 2, "n": 2},
    "S5": {"w": 0, "n": 1},
    "S6": {"w": 1, "n": 1},
}
def plot_tree_with_labels(edges=edges, nodes=nodes, root="S0", width=1.0, vert_gap=0.1, vert_loc=0, xcenter=1):
    def hierarchical_pos(G, root, width, vert_gap, vert_loc, xcenter):
        if not nx.is_tree(G):
            raise TypeError("Il grafo deve essere un albero!")

        def _hierarchy_pos(G, node, width, vert_gap, vert_loc, xcenter, pos=None, parent=None):
            if pos is None:
                pos = {node: (xcenter, vert_loc)}
            else:
                pos[node] = (xcenter, vert_loc)
            children = list(G.neighbors(node))
            if not isinstance(G, nx.DiGraph) and parent is not None:
                children.remove(parent)
            if len(children) != 0:
                dx = width / len(children)
                next_x = xcenter - width / 2 - dx / 2
                for child in children:
                    next_x += dx
                    pos = _hierarchy_pos(G, child, width=dx, vert_gap=vert_gap,
                                         vert_loc=vert_loc - vert_gap, xcenter=next_x, pos=pos, parent=node)
            return pos

        return _hierarchy_pos(G, root, width, vert_gap, vert_loc, xcenter)

    G = nx.DiGraph()
    for parent, child, label in edges:
        G.add_edge(parent, child, label=label)

    pos = hierarchical_pos(G, root, width, vert_gap, vert_loc, xcenter)

    node_labels = {
        node: f"{node}\nW={nodes[node]['w']}\nN={nodes[node]['n']}"
        for node in G.nodes()
    }

    plt.figure(figsize=(10, 8))
    nx.draw(G, pos, with_labels=False, arrows=True, node_size=2000, node_color='#ADD8E6')
    nx.draw_networkx_labels(G, pos, labels=node_labels, font_size=10, font_color='black')
    edge_labels = nx.get_edge_attributes(G, 'label')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    plt.title("MonteCarlo Tree Search")
    plt.axis("off")
    plt.show()

if __name__ == "__main__":
    plot_tree_with_labels()