Comparator of Kruskal's and Prim's MST searching algorithms

graph.txt

- 5 6 4 - -
5 - 1 2 - -
6 1 - 2 5 3
4 2 2 - - 4
- - 5 - - 4
- - 3 4 4 -

MinSpanTree.py

# MinSpanTree - Comparator of Kruskal's and Prim's MST searching algorithms
# Requires `graph.txt` - an adjacency matrix that defines the graph
# Put dashes (-) where there's no way

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import networkx as nx
from queue import PriorityQueue

def matrix_to_graph(matrix):
    oriented = False
    for i in range(len(matrix)):
        for j in range(len(matrix)):
            if matrix[i][j] != matrix[j][i]:
                oriented = True
                break
        if oriented:
            break

    graph = nx.DiGraph() if oriented else nx.Graph()
    for i, row in enumerate(matrix):
        for j, cell in enumerate(row):
            if cell != inf:
                graph.add_edge(i, j, weight=cell)
    return graph


class SetNode:
    def __init__(self, v):
        self.node = v
        self.parent = self
        self.rank = 0

    def root(self):
        curr = self
        if curr != curr.parent:
            curr.parent = curr.parent.root()
        return curr.parent

    def union(self, other):
        self_root = self.root()
        other_root = other.root()
        if self_root.node == other_root.node:
            return

        if self_root.rank > other_root.rank:
            other_root.parent = self_root
        else:
            self_root.parent = other_root
            if self_root.rank == other_root.rank:
                other_root.rank += 1

    def __repr__(self):
        form = 'Node {}, rank: {}, parent: {}'
        return form.format(self.node, self.rank, self.parent.node)


inf = float('inf')
with open('graph.txt') as f:
    g = [[int(i) if i != '-' else inf for i in line.split()] for line in f]

def kruskal(g):
    n = len(g)
    edges = []
    for i in range(n):
        for j in range(n):
            if g[i][j] != inf:
                edges.append((g[i][j], i, j))

    edges.sort(key = lambda x: x[0])

    verts = {i: SetNode(i) for i in range(n)}

    graph_edges = set()

    for ln, u, v in edges:
        if verts[u].root() != verts[v].root():
            verts[u].union(verts[v])
            graph_edges.add((u, v))
    return graph_edges


def modify(q, v, d):
    nq = PriorityQueue()
    while not q.empty():
        cost_i, i = q.get()
        if i != v:
            nq.put((cost_i, i))
        else:
            nq.put((d, v))
    return nq


def prim(g, s = 0):
    cost = [inf] * len(g)
    cost[s] = 0
    prev = [None] * len(g)

    q = PriorityQueue()
    verts = set(range(len(g)))
    in_q = set()

    for i in verts:
        q.put((cost[i], i))
        in_q.add(i)

    while not q.empty():
        cost_v, v = q.get()
        in_q.remove(v)
        for z in range(len(g)):
            if g[v][z] != inf:
                if z in in_q and cost[z] > g[v][z]:
                    cost[z] = g[v][z]
                    prev[z] = v
                    q = modify(q, z, cost[z])

    path = set()
    for u, v in enumerate(prev):
        if u is None or v is None:
            continue
        path.add((v, u))

    return path

all_verts = set(range(len(g)))
kr_edges = kruskal(g)
pr_edges = prim(g)

def join(edges1, edges2):
    both = set()
    for u, v in edges1:
        if (u, v) in edges2 or (v, u) in edges2:
            both.add((u, v))
    return both


G = matrix_to_graph(g)
pos = nx.circular_layout(G)

nodes = nx.draw_networkx_nodes(G, pos,
                               node_color = 'snow')
nodes.set_edgecolor('black')
nx.draw_networkx_labels(G, pos, {i: str(i) for i in all_verts}, font_size = 10)

nx.draw_networkx_edges(G, pos)
nx.draw_networkx_edges(G, pos, edgelist = kr_edges,
                       edge_color = 'lightblue')
nx.draw_networkx_edges(G, pos, edgelist = pr_edges,
                       edge_color = 'lightcoral')
nx.draw_networkx_edges(G, pos, edgelist = join(kr_edges, pr_edges),
                       edge_color = 'plum')
nx.draw_networkx_edge_labels(G, pos,
                             edge_labels = nx.get_edge_attributes(G, 'weight'))

blue = mpatches.Patch(color = 'lightblue', label = 'Kruskal\'s MST')
red = mpatches.Patch(color = 'lightcoral', label = 'Prim\'s MST')
purp = mpatches.Patch(color = 'plum', label = 'Matching edges')

plt.legend(handles = [blue, red, purp],
           bbox_transform = plt.gcf().transFigure,
           bbox_to_anchor = (1, 1))

plt.axis('off')
fig = plt.gcf()
fig.canvas.set_window_title('Search for a minimal spanning tree')
plt.show()