Reduce the complexity of graphs/minimum_spanning_tree_prims.py

.flake8

@@ -1,7 +1,7 @@
 [flake8]
 max-line-length = 88
 # max-complexity should be 10
-max-complexity = 21
+max-complexity = 20
 extend-ignore =
     # Formatting style for `black`
     E203  # Whitespace before ':'

graphs/minimum_spanning_tree_prims.py

@@ -2,17 +2,17 @@
 from collections import defaultdict
 
 
-def prisms_algorithm(l):  # noqa: E741
+class Heap:
+    def __init__(self):
+        self.node_position = []
 
-    node_position = []
+    def get_position(self, vertex):
+        return self.node_position[vertex]
 
-    def get_position(vertex):
-        return node_position[vertex]
+    def set_position(self, vertex, pos):
+        self.node_position[vertex] = pos
 
-    def set_position(vertex, pos):
-        node_position[vertex] = pos
-
-    def top_to_bottom(heap, start, size, positions):
+    def top_to_bottom(self, heap, start, size, positions):
         if start > size // 2 - 1:
             return
         else:
@@ -28,14 +28,14 @@ def top_to_bottom(heap, start, size, positions):
                 heap[m], positions[m] = heap[start], positions[start]
                 heap[start], positions[start] = temp, temp1
 
-                temp = get_position(positions[m])
-                set_position(positions[m], get_position(positions[start]))
-                set_position(positions[start], temp)
+                temp = self.get_position(positions[m])
+                self.set_position(positions[m], self.get_position(positions[start]))
+                self.set_position(positions[start], temp)
 
-                top_to_bottom(heap, m, size, positions)
+                self.top_to_bottom(heap, m, size, positions)
 
     # Update function if value of any node in min-heap decreases
-    def bottom_to_top(val, index, heap, position):
+    def bottom_to_top(self, val, index, heap, position):
         temp = position[index]
 
         while index != 0:
@@ -47,59 +47,78 @@ def bottom_to_top(val, index, heap, position):
             if val < heap[parent]:
                 heap[index] = heap[parent]
                 position[index] = position[parent]
-                set_position(position[parent], index)
+                self.set_position(position[parent], index)
             else:
                 heap[index] = val
                 position[index] = temp
-                set_position(temp, index)
+                self.set_position(temp, index)
                 break
             index = parent
         else:
             heap[0] = val
             position[0] = temp
-            set_position(temp, 0)
+            self.set_position(temp, 0)
 
-    def heapify(heap, positions):
+    def heapify(self, heap, positions):
         start = len(heap) // 2 - 1
         for i in range(start, -1, -1):
-            top_to_bottom(heap, i, len(heap), positions)
+            self.top_to_bottom(heap, i, len(heap), positions)
 
-    def delete_minimum(heap, positions):
+    def delete_minimum(self, heap, positions):
         temp = positions[0]
         heap[0] = sys.maxsize
-        top_to_bottom(heap, 0, len(heap), positions)
+        self.top_to_bottom(heap, 0, len(heap), positions)
         return temp
 
-    visited = [0 for i in range(len(l))]
-    nbr_tv = [-1 for i in range(len(l))]  # Neighboring Tree Vertex of selected vertex
+
+def prisms_algorithm(adjacency_list):
+    """
+    >>> adjacency_list = {0: [[1, 1], [3, 3]],
+    ...                   1: [[0, 1], [2, 6], [3, 5], [4, 1]],
+    ...                   2: [[1, 6], [4, 5], [5, 2]],
+    ...                   3: [[0, 3], [1, 5], [4, 1]],
+    ...                   4: [[1, 1], [2, 5], [3, 1], [5, 4]],
+    ...                   5: [[2, 2], [4, 4]]}
+    >>> prisms_algorithm(adjacency_list)
+    [(0, 1), (1, 4), (4, 3), (4, 5), (5, 2)]
+    """
+
+    heap = Heap()
+
+    visited = [0 for i in range(len(adjacency_list))]
+    nbr_tv = [
+        -1 for i in range(len(adjacency_list))
+    ]  # Neighboring Tree Vertex of selected vertex
     # Minimum Distance of explored vertex with neighboring vertex of partial tree
     # formed in graph
     distance_tv = []  # Heap of Distance of vertices from their neighboring vertex
     positions = []
 
-    for x in range(len(l)):
+    for x in range(len(adjacency_list)):
         p = sys.maxsize
         distance_tv.append(p)
         positions.append(x)
-        node_position.append(x)
+        heap.node_position.append(x)
 
     tree_edges = []
     visited[0] = 1
     distance_tv[0] = sys.maxsize
-    for x in l[0]:
+    for x in adjacency_list[0]:
         nbr_tv[x[0]] = 0
         distance_tv[x[0]] = x[1]
-    heapify(distance_tv, positions)
+    heap.heapify(distance_tv, positions)
 
-    for _ in range(1, len(l)):
-        vertex = delete_minimum(distance_tv, positions)
+    for _ in range(1, len(adjacency_list)):
+        vertex = heap.delete_minimum(distance_tv, positions)
         if visited[vertex] == 0:
             tree_edges.append((nbr_tv[vertex], vertex))
             visited[vertex] = 1
-            for v in l[vertex]:
-                if visited[v[0]] == 0 and v[1] < distance_tv[get_position(v[0])]:
-                    distance_tv[get_position(v[0])] = v[1]
-                    bottom_to_top(v[1], get_position(v[0]), distance_tv, positions)
+            for v in adjacency_list[vertex]:
+                if visited[v[0]] == 0 and v[1] < distance_tv[heap.get_position(v[0])]:
+                    distance_tv[heap.get_position(v[0])] = v[1]
+                    heap.bottom_to_top(
+                        v[1], heap.get_position(v[0]), distance_tv, positions
+                    )
                     nbr_tv[v[0]] = vertex
     return tree_edges
 
@@ -108,9 +127,9 @@ def delete_minimum(heap, positions):
     # < --------- Prims Algorithm --------- >
     n = int(input("Enter number of vertices: ").strip())
     e = int(input("Enter number of edges: ").strip())
-    adjlist = defaultdict(list)
+    adjacency_list = defaultdict(list)
     for x in range(e):
-        l = [int(x) for x in input().strip().split()]  # noqa: E741
-        adjlist[l[0]].append([l[1], l[2]])
-        adjlist[l[1]].append([l[0], l[2]])
-    print(prisms_algorithm(adjlist))
+        edge = [int(x) for x in input().strip().split()]
+        adjacency_list[edge[0]].append([edge[1], edge[2]])
+        adjacency_list[edge[1]].append([edge[0], edge[2]])
+    print(prisms_algorithm(adjacency_list))