minimum_spanning_tree_boruvka.py

classGraph:
"""
 Data structure to store graphs (based on adjacency lists)
 """

def__init__(self):
self.num_vertices=0
self.num_edges=0
self.adjacency= {}

defadd_vertex(self, vertex):
"""
 Adds a vertex to the graph

 """
ifvertexnotinself.adjacency:
self.adjacency[vertex] = {}
self.num_vertices+=1

defadd_edge(self, head, tail, weight):
"""
 Adds an edge to the graph

 """

self.add_vertex(head)
self.add_vertex(tail)

ifhead==tail:
return

self.adjacency[head][tail] =weight
self.adjacency[tail][head] =weight

defdistinct_weight(self):
"""
 For Boruvks's algorithm the weights should be distinct
 Converts the weights to be distinct

 """
edges=self.get_edges()
foredgeinedges:
head, tail, weight=edge
edges.remove((tail, head, weight))
foriinrange(len(edges)):
edges[i] =list(edges[i])

edges.sort(key=lambdae: e[2])
foriinrange(len(edges) -1):
ifedges[i][2] >=edges[i+1][2]:
edges[i+1][2] =edges[i][2] +1
foredgeinedges:
head, tail, weight=edge
self.adjacency[head][tail] =weight
self.adjacency[tail][head] =weight

def__str__(self):
"""
 Returns string representation of the graph
 """
string=""
fortailinself.adjacency:
forheadinself.adjacency[tail]:
weight=self.adjacency[head][tail]
string+=f"{head} -> {tail} == {weight}\n"
returnstring.rstrip("\n")

defget_edges(self):
"""
 Returna all edges in the graph
 """
output= []
fortailinself.adjacency:
forheadinself.adjacency[tail]:
output.append((tail, head, self.adjacency[head][tail]))
returnoutput

defget_vertices(self):
"""
 Returns all vertices in the graph
 """
returnself.adjacency.keys()

@staticmethod
defbuild(vertices=None, edges=None):
"""
 Builds a graph from the given set of vertices and edges

 """
g=Graph()
ifverticesisNone:
vertices= []
ifedgesisNone:
edge= []
forvertexinvertices:
g.add_vertex(vertex)
foredgeinedges:
g.add_edge(*edge)
returng

classUnionFind:
"""
 Disjoint set Union and Find for Boruvka's algorithm
 """

def__init__(self):
self.parent= {}
self.rank= {}

def__len__(self):
returnlen(self.parent)

defmake_set(self, item):
ifiteminself.parent:
returnself.find(item)

self.parent[item] =item
self.rank[item] =0
returnitem

deffind(self, item):
ifitemnotinself.parent:
returnself.make_set(item)
ifitem!=self.parent[item]:
self.parent[item] =self.find(self.parent[item])
returnself.parent[item]

defunion(self, item1, item2):
root1=self.find(item1)
root2=self.find(item2)

ifroot1==root2:
returnroot1

ifself.rank[root1] >self.rank[root2]:
self.parent[root2] =root1
returnroot1

ifself.rank[root1] <self.rank[root2]:
self.parent[root1] =root2
returnroot2

ifself.rank[root1] ==self.rank[root2]:
self.rank[root1] +=1
self.parent[root2] =root1
returnroot1
returnNone

@staticmethod
defboruvka_mst(graph):
"""
 Implementation of Boruvka's algorithm
 >>> g = Graph()
 >>> g = Graph.build([0, 1, 2, 3], [[0, 1, 1], [0, 2, 1],[2, 3, 1]])
 >>> g.distinct_weight()
 >>> bg = Graph.boruvka_mst(g)
 >>> print(bg)
 1 -> 0 == 1
 2 -> 0 == 2
 0 -> 1 == 1
 0 -> 2 == 2
 3 -> 2 == 3
 2 -> 3 == 3
 """
num_components=graph.num_vertices

union_find=Graph.UnionFind()
mst_edges= []
whilenum_components>1:
cheap_edge= {}
forvertexingraph.get_vertices():
cheap_edge[vertex] =-1

edges=graph.get_edges()
foredgeinedges:
head, tail, weight=edge
edges.remove((tail, head, weight))
foredgeinedges:
head, tail, weight=edge
set1=union_find.find(head)
set2=union_find.find(tail)
ifset1!=set2:
ifcheap_edge[set1] ==-1orcheap_edge[set1][2] >weight:
cheap_edge[set1] = [head, tail, weight]

ifcheap_edge[set2] ==-1orcheap_edge[set2][2] >weight:
cheap_edge[set2] = [head, tail, weight]
forhead_tail_weightincheap_edge.values():
ifhead_tail_weight!=-1:
head, tail, weight=head_tail_weight
ifunion_find.find(head) !=union_find.find(tail):
union_find.union(head, tail)
mst_edges.append(head_tail_weight)
num_components=num_components-1
mst=Graph.build(edges=mst_edges)
returnmst