import networkx as nx
import matplotlib.pyplot as plt


def kruskal(G, mx=False):
    sp = []
    n = len(G.nodes.items())
    for (u, v, wt) in G.edges.data('weight'):
        sp.append((u, v, wt))
    if not mx:
        sp.sort(key=lambda x: x[2])
        comp = [i for i in range(n)]
        ans = 0
        spp = {}
        x = []
        for start, end, weight in sp:
            start -= 1
            end -= 1
            if comp[start] != comp[end]:
                ans += weight
                a = comp[start]
                b = comp[end]
                for i in range(n):
                    if comp[i] == b:
                        comp[i] = a
                spp[(start + 1, end + 1)] = weight
                x.append((start + 1, end + 1))
        return spp, ans, x
    else:
        sp.sort(key=lambda x: -x[2])
        comp = [i for i in range(n)]
        ans = 0
        spp = {}
        x = []
        for start, end, weight in sp:
            start -= 1
            end -= 1
            if comp[start] != comp[end]:
                ans += weight
                a = comp[start]
                b = comp[end]
                for i in range(n):
                    if comp[i] == b:
                        comp[i] = a
                spp[(start + 1, end + 1)] = weight
                x.append((start + 1, end + 1))
        return spp, ans, x


G = nx.Graph()

G.add_node(1, pos=(0, 4))
G.add_node(2, pos=(2, 3))
G.add_node(3, pos=(4, 0))
G.add_node(4, pos=(2, -3))
G.add_node(5, pos=(0, -4))
G.add_node(6, pos=(-3, -2))
G.add_node(7, pos=(-4, 0))
G.add_node(8, pos=(-3, 2))

G.add_edge(1, 2, weight=2)
G.add_edge(1, 8, weight=1)

G.add_edge(2, 4, weight=1)
G.add_edge(2, 6, weight=3)
G.add_edge(2, 8, weight=3)

G.add_edge(3, 4, weight=2)
G.add_edge(3, 7, weight=6)

G.add_edge(4, 5, weight=2)
G.add_edge(4, 8, weight=3)

G.add_edge(5, 6, weight=1)

sp, ans, spp = kruskal(G)

pos = nx.get_node_attributes(G, 'pos')
labels = nx.get_edge_attributes(G, 'weight')

nx.draw(G, pos, with_labels=True, font_weight='bold', font_color="white", font_size=14)
nx.draw_networkx_nodes(G, pos, node_size=400, node_color="black")

nx.draw_networkx_edges(G, pos, edge_color="blue", edgelist=spp)
nx.draw_networkx_edge_labels(G, pos, edge_labels=labels, font_size=16)
nx.draw_networkx_edge_labels(G, pos, edge_labels=sp, font_color="blue", font_size=16)
print(f"Вес остава: {ans}")

plt.show()