myGraph.py 4.07 KB
import os
import re

class Graph:
	"""
	docstring for graph
	"""
	def __init__(self):
		"""
		use 0,1, 2, 3 as vertices
		and label the vertices using their variable name
		"""
		self.vertices = []
		self.edges = {}  #from key to value (parent -> child)
		self.parents = {} #the parents of the key, no label
		self.edge_labelling = {}
		self.vertice_labelling = {}


	def add_edge(self, start, end, label = None):
		if isinstance(start, int):
			start = [start]
			end = [end]
			label = [label] if label else None

		if label == None:

			for s, e in zip(start, end):
				self.edges[s].append(e)
				self.parents[e].append(s)
		else:

			for s, e, l in zip(start, end, label):
				self.edges[s].append(e)
				self.parents[e].append(s)
				self.edge_labelling[(s, e)] = l




	def add_vertice(self, vertice, label = None):
		if isinstance(vertice, int):
			vertice = [vertice]
			label = [label] if label else None


		if label == None:
			for v in vertice:
				self.vertices.append(v)
				self.edges[v] = []
				self.parents[v] = []
		else:
			for v, l in zip(vertice, label):
				self.vertices.append(v)
				self.vertice_labelling[v] = l
				self.edges[v] = []
				self.parents[v] = []

	def dot_to_graph(self, dot, d = None):
		"""
		txt is a dot txt file names, it is a string and 
		the location of txt is define in d, and the 
		default location is current working directory
		"""
		
		#load the dot data
		if not d:
			d = os.getcwd()
		f = '/'.join([d, dot])
		vertices = [line.rstrip('\n') for line in open(f) if '[label=' in line]
		edges = [line.rstrip('\n') for line in open(f) if '->' in line]

		#add vertices
		for l in vertices:
			l = l.replace(' ','')
			pattern = '([0-9]+)\[label="(.+)"\];'
			label = re.sub(pattern, '\\2',l)
			v = int(re.sub(pattern, '\\1', l))
			self.add_vertice(v, label)

		#add edges
		for l in edges:
			l = l.replace(' ', '')
			pattern = '([0-9]+)->([0-9]+);'
			s = int(re.sub(pattern, '\\1', l))
			e = int(re.sub(pattern, '\\2', l))
			self.add_edge(s, e)





	def graph_to_dot(self, name, label = None, labeljust = None, d = None):
		"""
		save a corresponding dot txt file in the 
		specific directory (define in d, and the 
		default location is current working directory)
		the dot txt file name is given by name, it 
		is a string

		"""
		if not d:
			d = os.getcwd()
		f = d + '/' + name + '.txt'
		tmp = open(f, 'w')
		tmp.write('digraph abstract {\n')

		#write label
		if not label:
			tmp.write('label = "'+label+'";\n')

		#Write labeljust
		if not labeljust:
			tmp.write('labeljust = "'+
				labeljust+'";\n')

		#write vertices
		for v, l in zip(self.vertices, self.vertice_labelling):
			tmp.write('\t' + v 
				+ ' [label="' + l + '"];\n')

		tmp.wirte('\n')

		#write edges
		all_edge = [(s, e) for e, v in self.parents for s in v]
		for s, e in all_edge:
			tmp.write('\t' + s + '->' + e + ';\n')

		tmp.write('\n')

		tmp.write('}\n')
		tmp.close()

		


	def delete_verrtice(self, vertice):
		pass

	def delete_edge(self, start, end):
		pass

	def ford_bellman(self, s):
		"""
		input s is a vertice
		output is a dict
		Bellman Ford algorithm to calculate the shortest distance
		"""
		D = {}
		inf = len(self.vertices)

		for x in self.vertices:
			if x in self.edges[s]:
				D[x] = 1
			else:
				D[x] = inf

		D[s] = 0

		for i in xrange(len(self.vertices) - 2):
			for v in self.vertices:
				if v != s:
					for u in self.vertices:
						if v in self.edges[u]:
							D[v] = min(D[v], D[u] + 1)

		return D


	def avg_dist(self):
		s = 0
		for v in self.vertices:
			s += sum(self.ford_bellman(v).values())
		return s * 1.0 / (len(self.vertices) ** 2 - len(self.vertices))

	def make_undirected(self):
		for v in self.vertices:
			for w in self.vertices:
				if (w in self.edges[v]) and (not v in self.edges[w]):
					self.edges[w].append(v)




	def fromParents(self, verts, pars):
		"""
		add edges whose parents have been given
		"""

		self.vertices = []
		for v in verts:
			self.add_vertice(v)
			self.edges[v] = []

		for v, par in zip(verts, pars):
			for p in par:
				self.add_edge(p, v)


	def num_edges(self):
		return sum(map(len, self.edges.values()))