最小生成树：距离和图形

Question

我遇到了最小的生成树问题 我找到了一个解决问题的代码并返回解决方案的图表，我已对其进行了一些编辑，但我找不到一种方法将距离返回到图形，就像我一样没有太多的蟒蛇经验 你能帮助我吗 ？感谢

    import numpy as np
    from scipy.spatial.distance import pdist, squareform
    import matplotlib.pyplot as plt


def minimum_spanning_tree(X, copy_X=True):
    if copy_X:
        X = X.copy()

    if X.shape[0] != X.shape[1]:
        raise ValueError("X needs to be square matrix of edge weights")
    n_vertices = X.shape[0]
    spanning_edges = []

    visited_vertices = [0]
    num_visited = 1

    diag_indices = np.arange(n_vertices)
    X[diag_indices, diag_indices] = np.inf

    while num_visited != n_vertices:
        new_edge = np.argmin(X[visited_vertices], axis=None)
        new_edge = divmod(new_edge, n_vertices)
        new_edge = [visited_vertices[new_edge[0]], new_edge[1]]
        spanning_edges.append(new_edge)
        visited_vertices.append(new_edge[1])
        X[visited_vertices, new_edge[1]] = np.inf
        X[new_edge[1], visited_vertices] = np.inf
        num_visited += 1
    return np.vstack(spanning_edges)


def test_mst():
    n = int(input())
    P = np.array([], int)
    for i in range(0,n):
        y = list(map(int,input().split()))
        P = np.append(P,[y[0],y[1]], axis=0)
    P = P.reshape(n,2)
    X = squareform(pdist(P))
    edge_list = minimum_spanning_tree(X)
    plt.scatter(P[:, 0], P[:, 1])

    for edge in edge_list:
        i, j = edge
        plt.plot([P[i, 0], P[j, 0]], [P[i, 1], P[j, 1]], c='r')
    plt.show()

if __name__ == "__main__":
    test_mst()