Graphs in Python

graph_min_span_tree.py

#!/usr/bin/python


from __future__ import print_function


class Vertex:
	
	def __init__(self, label, value=None):
		self.label = label
		self.value = value

	def __repr__(self):
		return str(self.label) + ':' + str(self.value) + ' '

	def __str__(self):
		return self.__repr__()


class Edge:

	def __init__(self, vertex1, vertex2, weight=None, directed=False):
		self.vertex1 = vertex1
		self.vertex2 = vertex2
		self.directed = directed
		self.weight = weight

	def __repr__(self):
		a = '' if self.weight==None else str(self.weight)
		if self.directed:
			return str(self.vertex1.__repr__()) + '--'+a+'--> ' + str(self.vertex2.__repr__()) + ' '
		else:
			return str(self.vertex1.__repr__()) + '--'+a+'-- ' + str(self.vertex2.__repr__()) + ' '

	def __str__(self):
		return self.__repr__()

	def set_weight(weight):
		self.weight = weight
	
	def make_directed():
		self.directed = True; 

	def get_weight():
		return self.weight
		

class Graph:

	def __init__(self, set_of_vertices, set_of_edges):
		self.vertices = set_of_vertices
		self.edges = set_of_edges

	def copy(self):
		return Graph(self.vertices, self.edges)
	
	def add_edge(self, edge):
		#Check if 'edge' is a set of edges or just an edge
		if type(edge) is set:
			self.edges = self.edges.union(edge)
		elif isinstance(edge, Edge):
			self.edges = self.edges.union(set([edge]))
		else:
			raise Exception("Illegal Type On Argument 1 of add_edge() method")		


	def add_vertex(self, vertex):
		#Check if 'vertex' is a set of vertices or just a vertex
		if type(vertex) is set:
			self.vertices = self.vertices.union(vertex)
		elif isinstance(vertex, Vertex):
			self.vertices = self.vertices.union(set([vertex]))
		else:
			raise Exception("Illegal Type On Argument 1 of add_vertex() method")

	def get_vertices(self):
		return self.vertices
	
	def get_edges(self):
		return self.edges


	def debug(self):
		print("Vertices: " , self.vertices)
		print("Edges: ", self.edges)

	def visualise(self):
		#Will Visualise the graph
		vertices_used = set() #Empty set

	def toMinSpanTree(self, startingVertex=False):
		#Use Prim's Algorithm to create a new graph which is a Minimum Spanning Tree
		
		#Start with a set of new Vertex
		if len(self.vertices) == 0 :
			raise Exception("Cannot Convert Graph to Spanning Tree. No vertices")
		if len(self.edges) == 0:
			#No edges, means this is a MST forest with no connected vertices, return 
			return self.copy()

		if not startingVertex:
			#Get "probably" the first vertex in the set
			for elem in self.vertices:
				break;
		else:
			elem = startingVertex
			
		V_new = set([elem])
		E_new = set()
		
		#To take a reference weight for comparision
		#for weight in self.edges:
		#	break
		#weight = weight.weight
		
		while(V_new != self.vertices):
			count = 0
			new_edge = {}
			for edge in self.edges:
				if ((edge.vertex1 in V_new) and (edge.vertex2 not in V_new)):
					new_edge = edge if count==0 else (new_edge if (edge.weight > new_edge.weight) else edge)
					count = count + 1
			#Min. Edge found ?
			#Hmm..Probably.
			#In that case:  add that to set
			if type(new_edge) is not dict:
				E_new.add(new_edge)
				V_new.add(new_edge.vertex2)
			else:
				for edge in self.edges:
					if ((edge.vertex2 in V_new) and (edge.vertex1 not in V_new)):
						new_edge = edge if count==0 else (new_edge if (edge.weight > new_edge.weight) else edge)
						count = count + 1
				#Min Edge Found?
				#Hmm..Probably.
				#In that case:  add that to set
				if type(new_edge) is not dict:
					E_new.add(new_edge)
					V_new.add(new_edge.vertex1)
				else:
					#No spanning possible : I'm not sure if this ever reaches!
					debugText = "Cannot Span Starting from Vertex: "+ str(elem) + "\n--DEBUG_INFO--" + "\nSet Of Used Vertices: "+ str(V_new) + "\nSet Of new Edges: "+ str(E_new) + "\n-------"

					return type("", 
						   (),
						   dict
						   (debug = lambda self: print(debugText) ))()
		return Graph(V_new, E_new) #or V, E_new

vertexA = Vertex('A',1)
vertexB = Vertex('B',2)
vertexC = Vertex('C',3)
vertexD = Vertex('D',4)

edge1  = Edge(vertexA, vertexB, 2);
edge2  = Edge(vertexA, vertexD, 1);
edge3  = Edge(vertexB, vertexD, 2);
edge4  = Edge(vertexD, vertexC, 3);

vertices = set([vertexA, vertexB, vertexC, vertexD])
edges = set([edge1, edge2, edge3, edge4])

graph_initial = Graph(vertices, edges)
mst_1 = graph_initial.toMinSpanTree(vertexB)
mst_2 = graph_initial.toMinSpanTree(vertexA)
print()
print()
print("Initial Graph: ")
print()
graph_initial.debug()
print()
print()
print("Minimum Span Tree Derived from Initial Graph using Prim's Algorithm. (Starting Vertex: Vertex B)")
print()
mst_1.debug()
print()
print()
print("Minimum Span Tree Derived from Initial Graph using Prim's Algorithm. (Starting Vertex: Vertex A)")
print()
mst_2.debug()

graphs_basic.py

#!/usr/bin/python

class Vertex:
	
	def __init__(self, label, value=None):
		self.label = label
		self.value = value

	def __repr__(self):
		return str(self.label) + ':' + str(self.value) + ' '


class Edge:

	def __init__(self, vertex1, vertex2):
		self.vertex1 = vertex1
		self.vertex2 = vertex2
		self.directed = False #Default

	def __repr__(self):
		if self.directed:
			return str(self.vertex1.__repr__()) + '----> ' + str(self.vertex2.__repr__()) + ' '
		else:
			return str(self.vertex1.__repr__()) + '---- ' + str(self.vertex2.__repr__()) + ' '
		

class graph:

	def __init__(self, set_of_vertices, set_of_edges):
		self.vertices = set_of_vertices
		self.edges = set_of_edges
	
	def add_edge(self, edge):
		#Check if 'edge' is a set of edges or just an edge
		if type(edge) is set:
			self.edges = self.edges.union(edge)
		elif isinstance(edge, Edge):
			self.edges = self.edges.union(set([edge]))
		else:
			raise Exception("Illegal Type On Argument 1 of add_edge() method")
		


	def add_vertex(self, vertex):
		#Check if 'vertex' is a set of vertices or just a vertex
		if type(vertex) is set:
			self.vertices = self.vertices.union(vertex)
		elif isinstance(vertex, Vertex):
			self.vertices = self.vertices.union(set([vertex]))
		else:
			raise Exception("Illegal Type On Argument 1 of add_vertex() method")


	def debug(self):
		print "Vertices: " , self.vertices
		print "Edges: ", self.edges

		


#Example1: Construct a Graph: C4 
#Example2: Construct a Graph: K4

#Example1:
	#Create 4 Vertices

VertexA = Vertex('A',1);
VertexB = Vertex('B',2);
VertexC = Vertex('C',3);
VertexD = Vertex('D',4);

	#Create 4 Edges
Edge1 = Edge(VertexA,VertexB);
Edge2 = Edge(VertexB,VertexC);
Edge3 = Edge(VertexC,VertexD);
Edge4 = Edge(VertexD,VertexA);

	#Create Sets

vertices = set([VertexA, VertexB, VertexC, VertexD])
edges = set([Edge1, Edge2, Edge3, Edge4])

graph_example1 = graph(vertices,edges)
graph_example2 = graph(vertices,edges)
#Add Edge to graph_example2


VertexE = Vertex('E',5)

graph_example2.add_vertex(VertexE)

Edge6 = Edge(VertexA, VertexE)
Edge7 = Edge(VertexB, VertexE)
Edge8 = Edge(VertexC, VertexE)
Edge9 = Edge(VertexD, VertexE)

graph_example2.add_edge(set([Edge6,Edge7,Edge8,Edge9]))

#Debug (or rather dump)
graph_example1.debug()
graph_example2.debug()