From 415c9f5e6547457eb3546b467283cbd9e82e4eec Mon Sep 17 00:00:00 2001
From: Bruno Santos <7022432+dunderbruno@users.noreply.github.com>
Date: Sun, 1 Dec 2019 02:13:28 -0300
Subject: [PATCH] Improve prim.py (#1226)

* suiting PEP8

* create auxiliary function

* running example

* updating DIRECTORY.md
---
 DIRECTORY.md   |  3 ++
 graphs/prim.py | 76 +++++++++++++++++++++++++++++++++-----------------
 2 files changed, 53 insertions(+), 26 deletions(-)

diff --git a/DIRECTORY.md b/DIRECTORY.md
index 74c63d144..a3db3636e 100644
--- a/DIRECTORY.md
+++ b/DIRECTORY.md
@@ -82,6 +82,9 @@
     * [Red Black Tree](https://github.com/TheAlgorithms/Python/blob/master/data_structures/binary_tree/red_black_tree.py)
     * [Segment Tree](https://github.com/TheAlgorithms/Python/blob/master/data_structures/binary_tree/segment_tree.py)
     * [Treap](https://github.com/TheAlgorithms/Python/blob/master/data_structures/binary_tree/treap.py)
+  * Data Structures
+    * Heap
+      * [Heap Generic](https://github.com/TheAlgorithms/Python/blob/master/data_structures/data_structures/heap/heap_generic.py)
   * Disjoint Set
     * [Disjoint Set](https://github.com/TheAlgorithms/Python/blob/master/data_structures/disjoint_set/disjoint_set.py)
   * Hashing
diff --git a/graphs/prim.py b/graphs/prim.py
index 336424d2c..16cfaee08 100644
--- a/graphs/prim.py
+++ b/graphs/prim.py
@@ -2,26 +2,12 @@
 Prim's Algorithm.
 
 Determines the minimum spanning tree(MST) of a graph using the Prim's Algorithm
-
-Create a list to store x the vertices.
-G = [vertex(n) for n in range(x)]
-
-For each vertex in G, add the neighbors:
-G[x].addNeighbor(G[y])
-G[y].addNeighbor(G[x])
-
-For each vertex in G, add the edges:
-G[x].addEdge(G[y], w)
-G[y].addEdge(G[x], w)
-
-To solve run:
-MST = prim(G, G[0])
 """
 
 import math
 
 
-class vertex:
+class Vertex:
     """Class Vertex."""
 
     def __init__(self, id):
@@ -36,7 +22,7 @@ class vertex:
         self.key = None
         self.pi = None
         self.neighbors = []
-        self.edges = {}  # [vertex:distance]
+        self.edges = {}  # {vertex:distance}
 
     def __lt__(self, other):
         """Comparison rule to < operator."""
@@ -46,34 +32,72 @@ class vertex:
         """Return the vertex id."""
         return self.id
 
-    def addNeighbor(self, vertex):
+    def add_neighbor(self, vertex):
         """Add a pointer to a vertex at neighbor's list."""
         self.neighbors.append(vertex)
 
-    def addEdge(self, vertex, weight):
+    def add_edge(self, vertex, weight):
         """Destination vertex and weight."""
         self.edges[vertex.id] = weight
 
 
+def connect(graph, a, b, edge):
+    # add the neighbors:
+    graph[a - 1].add_neighbor(graph[b - 1])
+    graph[b - 1].add_neighbor(graph[a - 1])
+    # add the edges:
+    graph[a - 1].add_edge(graph[b - 1], edge)
+    graph[b - 1].add_edge(graph[a - 1], edge)
+
+
 def prim(graph, root):
     """
     Prim's Algorithm.
     Return a list with the edges of a Minimum Spanning Tree
     prim(graph, graph[0])
     """
-    A = []
+    a = []
     for u in graph:
         u.key = math.inf
         u.pi = None
     root.key = 0
-    Q = graph[:]
-    while Q:
-        u = min(Q)
-        Q.remove(u)
+    q = graph[:]
+    while q:
+        u = min(q)
+        q.remove(u)
         for v in u.neighbors:
-            if (v in Q) and (u.edges[v.id] < v.key):
+            if (v in q) and (u.edges[v.id] < v.key):
                 v.pi = u
                 v.key = u.edges[v.id]
     for i in range(1, len(graph)):
-        A.append([graph[i].id, graph[i].pi.id])
-    return A
+        a.append((int(graph[i].id) + 1, int(graph[i].pi.id) + 1))
+    return a
+
+
+def test_vector() -> None:
+    """
+    # Creates a list to store x vertices.
+    >>> x = 5
+    >>> G = [Vertex(n) for n in range(x)]
+
+    >>> connect(G, 1, 2, 15)
+    >>> connect(G, 1, 3, 12)
+    >>> connect(G, 2, 4, 13)
+    >>> connect(G, 2, 5, 5)
+    >>> connect(G, 3, 2, 6)
+    >>> connect(G, 3, 4, 6)
+    >>> connect(G, 0, 0, 0)  # Generate the minimum spanning tree:
+    >>> MST = prim(G, G[0])
+    >>> for i in MST:
+    ...     print(i)
+    (2, 3)
+    (3, 1)
+    (4, 3)
+    (5, 2)
+    """
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()