#!/usr/bin/env python
# coding: utf-8

# In[23]:


import numpy as np
from heapq import heappop, heappush


# In[24]:


## Graph


# In[25]:


V = ["A","B","C","D","E","F","G","H","I"]
E = {("A","B"):4, ("A","H"):8, ("B","H"): 11, ("B","C"): 8, ("H","I"): 7,("H","G"):1,("C","I"):2,("C","D"):7,("C","F"):4,("I","G"):6}
E[("D","E")] = 9
E[("F","E")] = 10
E[("D","F")] = 14
E[("G","F")] = 2


# In[26]:


def lightest_edge(E, visited_vertices):
    light = np.inf
    for i in E:
        if (i[0] in visited_vertices and not i[1] in visited_vertices): 
            if E[i] < light:
                light = E[i]
                aresta = i
                vertice = i[1]
        elif (i[1] in visited_vertices and not i[0] in visited_vertices):
            if E[i] < light:
                light = E[i]
                aresta = i
                vertice = i[0]
    return aresta,vertice


# In[27]:


def slow_prim(E,V):
    s = np.random.choice(V)
    MST = set()
    visited_vertices = {s}
    while len(visited_vertices) < len(V):
        aresta,v = lightest_edge(E, visited_vertices) #procura nas arestas em que x está e v não mas visitadas
        MST.add(aresta)
        visited_vertices.add(v)
    return MST


# In[28]:


slow_prim(E,V)


# In[29]:


## Grapghs with values in the vertices


# In[30]:


def add_edge(v, u, w):
    graph[v].append((u,w))
    graph[u].append((v,w)) # considera que o grafo eh nao direcionado


# In[34]:


graph = [[] for x in range(9)]
add_edge(0,1,4)
add_edge(0,7,8)
add_edge(1,7,11)
add_edge(1,2,8)
add_edge(2,8,2)
add_edge(2,3,7)
add_edge(2,5,4)
add_edge(3,4,9)
add_edge(6,7,1)
add_edge(6,8,6)
add_edge(7,8,7)
add_edge(4,5,10)
add_edge(3,5,14)
add_edge(5,6,2)
#Graph = dict()
#Letters = ["A","B","C","D","E","F","G","H","I"]
#for i in range(len(graph)):
 #   Graph[Letters[i]] = graph[i]
  #  for j in range(len(graph[i])):
   #     Graph[Letters[i]][j] = (Letters[graph[i][j][0]],graph[i][j][1])


# In[35]:


def prim(graph):
    
    edge = [(-1,-1,-1)] * len(graph)    #weight of edge, predecessor, node
    visited = [False] * len(graph)
    
    inicial = np.random.randint(9)
    edge[inicial] = (0, -1, inicial)
    heap = [(0,inicial)]
    
    while visited.count(True) < len(graph):
        v = -1
        while len(heap) > 0 and (v < 0 or visited[v]):
            v = heappop(heap)[1]
        visited[v] = True
        
        for (u,w) in graph[v]:
            if (edge[u][0] < 0 or edge[u][0] > w) and not visited[u] :
                edge[u] = (w,v,u)
                heappush(heap,(edge[u][0],u))
                
    V = ["A","B","C","D","E","F","G","H","I","Null"]
    for e in range(len(edge)):
        edge[e] = (V[edge[e][2]],V[edge[e][1]])        
    return edge #print de node and its sucessor


# In[36]:


prim(graph)


# In[ ]: