Clase Grafo con algunos metodos adicionales para el TP

clase_grafo_TP.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):
		""" Imprime 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

	def vertices_iniciales(self):
		# Devuelve el conjunto de vertices iniciales (o sea, vértices tales que d_in(v)=0 )
		hijos = set()
		for v in self.vertices:
			hijos.update(self.vecinos(v))
		return self.vertices.difference(hijos)

	def eliminar_arista(self, arista):
		# Elimina la arista (u,v). 'arista' debe ser una tupla o una lista de dos nodos.
		u, v = arista
		for vecino_con_peso in self.vecinos_con_peso(u):
			if vecino_con_peso[0] == v:
				self.lista_adyacencia[u].remove(vecino_con_peso)
				break

	def cantidad_de_aristas(self):
		return sum(len(self.vecinos(u)) for u in self.vertices)

	def peso_de_arista(self, arista):
		# Devuelve el peso de la arista (u,v). 'arista' debe ser una tupla o una lista de dos nodos.
		u, v = arista
		for vecino_con_peso in self.vecinos_con_peso(u):
			if vecino_con_peso[0] == v:
				return vecino_con_peso[1]
		print(f'La arista {arista} no forma parte del grafo')