Clase Grafo para la materia Investigacion Operativa

clase_grafo.py

"""
Clase Grafo para la materia Investigación Operativa, del Departamento de Matemática, Facultad de Ciencias Exactas y
Naturales, Universidad de Buenos Aires.
"""


class Grafo:

	def __init__(self):
		""" Crea un grafo vacío. Representación de un grafo por lista de adyacencia """
		self.lista_adyacencia = {}		# Diccionario
		self.vertices = set()
		
	def agregar_vertice(self, vertice):
		""" Agrega un vértice al grafo """
		self.vertices.add(vertice)
		
	def agregar_arista(self, arista):
		"""
		Se utiliza para agregar aristas. Si es un grafo pesado, el input debe ser: (u,v, peso)
		Si no es un grafo pesado, el input debe ser: (u, v).
		Si no es un grafo dirigido, agregar (u,v) y (v, u).
		Si u y/o v no están entre los vértices del grafo, los agrega.
		"""

		peso = 1
		if len(arista) == 3:
			salida, llegada, peso = arista
		elif len(arista) == 2:
			salida, llegada = arista
		else:
			print("Uso invalido. Ingresar [salida,llegada,peso]")

		# Si los vertices no pertenecian al grafo, se los agrega
		self.vertices.add(salida)
		self.vertices.add(llegada)
		if salida not in self.lista_adyacencia:
			self.lista_adyacencia[salida] = set()
		self.lista_adyacencia[salida].add((llegada, peso))
			
	def agregar_lista_aristas(self,lista_aristas):
		"""
		Permite agregar una lista de aristas, que debe venir dada como una lista de tuplas aptas para ser utilizadas
		con la funcion agregar_arista
		"""
		for arista in lista_aristas:
			self.agregar_arista(arista)
			
	def cantidad_de_vertices(self):
		return len(self.vertices)

	def vecinos(self, vertice):
		""" Devuelve un set con los vértices vecinos """
		if vertice not in self.lista_adyacencia:
			return set()
		else:
			return {u[0] for u in self.lista_adyacencia[vertice]}

	def vecinos_con_peso(self, vertice):
		""" Para grafos dirigidos. Devuelve un set con los vecios del vértice y sus correspondientes pesos. """
		if vertice not in self.lista_adyacencia:
			return set()
		else:
			return self.lista_adyacencia[vertice]

	def __str__(self):
		""" Se utiliza para imprimir el grafo """
		vertices = f'Vertices: {self.vertices}'
		ls_adyacencia = f'Lista de vecinos: {self.lista_adyacencia}'
		aristas = f'Aristas: {[(u, v, peso) for u, ady in self.lista_adyacencia.items() for v, peso in ady]}'
		return vertices + '\n' + ls_adyacencia + '\n' + aristas