from copy import deepcopy

class Grafo:
	def __init__(self, v = []):
		self.vertices = v
		self.arestas = {}
		self.distancias = {}

		if not v == []:
			for i in range(len(v)):
				self.arestas[v[i]] = []

	def insere_vertice(self, valor):
		self.vertices.append(valor)

	def insere_aresta(self, v1, v2, custo = 0):
		if not self.arestas[v1]:
			self.arestas[v1] = []
		if not self.arestas[v2]:
			self.arestas[v2] = []

		self.arestas[v1].append(v2)
		self.arestas[v2].append(v1)
		self.arestas[v1].sort()
		self.arestas[v2].sort()
		self.distancias[(v1, v2)] = custo

	def retira_vertice(self, v):
		if v in self.vertices:
			self.vertices.remove(v)
			del self.arestas[v]

		for vertice in self.vertices:
			if v in self.arestas[vertice]:
				self.arestas[vertice].remove(v)

		for aresta in deepcopy(self.distancias):
			if v in aresta:
				del self.distancias[aresta]

	def retira_aresta(self, v1, v2):
		if v1 not in self.vertices or v2 not in self.vertices: return
		if v1 in self.arestas[v2]:
			self.arestas[v2].remove(v1)
		if v2 in self.arestas[v1]:
			self.arestas[v1].remove(v2)

		for aresta in deepcopy(self.distancias):
			if aresta == (v1, v2) or aresta == (v2, v1):
				del self.distancias[aresta]


	# DIJKSTRA: menores caminhos de um vértice até todos os outros do grafo
	def dijkstra(self, ini):
		if ini not in self.vertices: return

		Q = []  # vértices não-visitados
		dist = {}
		prev = {}

		for v in self.vertices:
			dist[v] = float('inf')
			prev[v] = None
			Q.append(v)

		dist[ini] = 0

		while Q:
			u = Q[0]
			for x in Q:
				if dist[x] < dist[u]:
					u = x

			Q.remove(u)

			for v in self.arestas[u]:
				alt = dist[u]

				if (u,v) in self.distancias:
					alt += self.distancias[(u,v)]
				elif (v,u) in self.distancias:
					alt += self.distancias[(v,u)]

				if alt < dist[v]:
					dist[v] = alt
					prev[v] = u

		return dist # ou return dist[x] para retornar a distância de ini até x