{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from heapq import heappop, heappush"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = [\"A\",\"B\",\"C\",\"D\",\"E\",\"F\",\"G\",\"H\",\"I\"]\n",
    "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}\n",
    "E[(\"D\",\"E\")] = 9\n",
    "E[(\"F\",\"E\")] = 10\n",
    "E[(\"D\",\"F\")] = 14\n",
    "E[(\"G\",\"F\")] = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lightest_edge(E, visited_vertices):\n",
    "    light = np.inf\n",
    "    for i in E:\n",
    "        if (i[0] in visited_vertices and not i[1] in visited_vertices): \n",
    "            if E[i] < light:\n",
    "                light = E[i]\n",
    "                aresta = i\n",
    "                vertice = i[1]\n",
    "        elif (i[1] in visited_vertices and not i[0] in visited_vertices):\n",
    "            if E[i] < light:\n",
    "                light = E[i]\n",
    "                aresta = i\n",
    "                vertice = i[0]\n",
    "    return aresta,vertice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def slow_prim(E,V):\n",
    "    s = np.random.choice(V)\n",
    "    MST = set()\n",
    "    visited_vertices = {s}\n",
    "    while len(visited_vertices) < len(V):\n",
    "        aresta,v = lightest_edge(E, visited_vertices) #procura nas arestas em que x está e v não mas visitadas\n",
    "        MST.add(aresta)\n",
    "        visited_vertices.add(v)\n",
    "    return MST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('A', 'B'),\n",
       " ('A', 'H'),\n",
       " ('C', 'D'),\n",
       " ('C', 'F'),\n",
       " ('C', 'I'),\n",
       " ('D', 'E'),\n",
       " ('G', 'F'),\n",
       " ('H', 'G')}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "slow_prim(E,V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Grapghs with values in the vertices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_edge(v, u, w):\n",
    "    graph[v].append((u,w))\n",
    "    graph[u].append((v,w)) # considera que o grafo eh nao direcionado"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = [[] for x in range(9)]\n",
    "add_edge(0,1,4)\n",
    "add_edge(0,7,8)\n",
    "add_edge(1,7,11)\n",
    "add_edge(1,2,8)\n",
    "add_edge(2,8,2)\n",
    "add_edge(2,3,7)\n",
    "add_edge(2,5,4)\n",
    "add_edge(3,4,9)\n",
    "add_edge(6,7,1)\n",
    "add_edge(6,8,6)\n",
    "add_edge(7,8,7)\n",
    "add_edge(4,5,10)\n",
    "add_edge(3,5,14)\n",
    "add_edge(5,6,2)\n",
    "#Graph = dict()\n",
    "#Letters = [\"A\",\"B\",\"C\",\"D\",\"E\",\"F\",\"G\",\"H\",\"I\"]\n",
    "#for i in range(len(graph)):\n",
    " #   Graph[Letters[i]] = graph[i]\n",
    "  #  for j in range(len(graph[i])):\n",
    "   #     Graph[Letters[i]][j] = (Letters[graph[i][j][0]],graph[i][j][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prim(graph):\n",
    "    \n",
    "    edge = [(-1,-1,-1)] * len(graph)    #weight of edge, predecessor, node\n",
    "    visited = [False] * len(graph)\n",
    "    \n",
    "    inicial = np.random.randint(9)\n",
    "    edge[inicial] = (0, -1, inicial)\n",
    "    heap = [(0,inicial)]\n",
    "    \n",
    "    while visited.count(True) < len(graph):\n",
    "        v = -1\n",
    "        while len(heap) > 0 and (v < 0 or visited[v]):\n",
    "            v = heappop(heap)[1]\n",
    "        visited[v] = True\n",
    "        \n",
    "        for (u,w) in graph[v]:\n",
    "            if (edge[u][0] < 0 or edge[u][0] > w) and not visited[u] :\n",
    "                edge[u] = (w,v,u)\n",
    "                heappush(heap,(edge[u][0],u))\n",
    "                \n",
    "    V = [\"A\",\"B\",\"C\",\"D\",\"E\",\"F\",\"G\",\"H\",\"I\",\"Null\"]\n",
    "    for e in range(len(edge)):\n",
    "        edge[e] = (V[edge[e][2]],V[edge[e][1]])        \n",
    "    return edge #print de node and its sucessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('A', 'B'),\n",
       " ('B', 'Null'),\n",
       " ('C', 'B'),\n",
       " ('D', 'C'),\n",
       " ('E', 'D'),\n",
       " ('F', 'C'),\n",
       " ('G', 'F'),\n",
       " ('H', 'G'),\n",
       " ('I', 'C')]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prim(graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}