OmkarPathak · OmkarPathak · Aug 16, 2017 · Aug 14, 2017 · Aug 14, 2017 · Aug 14, 2017
diff --git a/pygorithm/data_structures/graph.py b/pygorithm/data_structures/graph.py
@@ -1,40 +1,99 @@
 # Author: OMKAR PATHAK
 # Created On: 12th August 2017
+from collections import defaultdict
+
 
 class Graph(object):
     def __init__(self):
-        self.graph = {}
+        self.graph = defaultdict(list)
         self.count = 0
 
     def print_graph(self):
         ''' for printing the contents of the graph '''
-        for i  in self.graph:
-            print(i,'->',' -> '.join([str(j) for j in self.graph[i]]))
+        for i in self.graph:
+            print(i, '->', ' -> '.join([str(j) for j in self.graph[i]]))
 
     def add_edge(self, from_vertex, to_vertex):
         ''' function to add an edge in the graph '''
         # check if vertex is already present
-        if from_vertex in self.graph.keys():
-            self.graph[from_vertex].append(to_vertex)
-            self.count += 1
-        else:
-            self.graph[from_vertex] = [to_vertex]
-            self.graph[to_vertex] = []
-            self.count += 1
+        self.graph[from_vertex].append(to_vertex)
+        self.count += 1
 
     def get_code(self):
-        ''' returns the code for the current class '''
+        """ returns the code for the current class """
         import inspect
         return inspect.getsource(Graph)
 
+
+class WeightedGraph:
+    """
+    A graph with a numerical value (weight) on edges
+    """
+    def __init__(self):
+        self.edges_weighted = []
+        self.vertexes = set()
+
+    def add_edge(self, u, v, weight):
+        """
+        :param u: from vertex - type : integer
+        :param v: to vertex - type : integer
+        :param weight: weight of the edge - type : numeric
+        """
+        edge = ((u, v), weight)
+        self.edges_weighted.append(edge)
+        self.vertexes.update((u, v))
+
+    def print_graph(self):
+        for (u, v), weight in self.edges_weighted:
+            print("%d -> %d weight: %d" % (u, v, weight))
+
+    def _set_of(self, vertex):
+        for tree in self.forest:
+            if vertex in tree:
+                return tree
+        return None
+
+    def _union(self, u_set, v_set):
+        self.forest.remove(u_set)
+        self.forest.remove(v_set)
+        self.forest.append(v_set + u_set)
+
+    def kruskal_mst(self):
+        """
+        Kruskal algorithm for finding the minimum spanning tree of a weighted graph.
+        This version use a union-find data structure.
+        More detailed info here: https://en.wikipedia.org/wiki/Kruskal%27s_algorithm
+        Author: Michele De Vita <mik3dev@gmail.com>
+        """
+        # sort by weight
+        self.edges_weighted.sort(key=lambda pair: pair[1])
+        edges_explored = []
+        self.forest = [[v] for v in self.vertexes]
+        for (u, v), weight in self.edges_weighted:
+            u_set, v_set = self._set_of(u), self._set_of(v)
+            if u_set != v_set:
+                self._union(u_set, v_set)
+                edges_explored.append(((u, v), weight))
+        return edges_explored
+
+    @staticmethod
+    def kruskal_complexity():
+        return '''Worst case: O(E log(V)) where E in the number of edges and V the number of vertexes'''
+
+    @classmethod
+    def kruskal_code(cls):
+        import inspect
+        return inspect.getsource(cls.kruskal_mst)
+
+
 class TopologicalSort(Graph):
     def topological_sort(self):
         ''' function for sorting graph elements using topological sort '''
-        visited = [False] * self.count           # Marking all vertices as not visited
-        stack = []                               # Stack for storing the vertex
+        visited = [False] * self.count  # Marking all vertices as not visited
+        stack = []  # Stack for storing the vertex
         for vertex in range(self.count):
             # Call the recursive function only if not visited
-            if visited[vertex] == False:
+            if not visited[vertex]:
                 self.topological_sort_rec(vertex, visited, stack)
 
         return stack
@@ -53,22 +112,23 @@ def topological_sort_rec(self, vertex, visited, stack):
             return
 
         # Push current vertex to stack which stores the result
-        stack.insert(0,vertex)
+        stack.insert(0, vertex)
 
     def get_code(self):
         ''' returns the code for the current class '''
         import inspect
         return inspect.getsource(TopologicalSort)
 
+
 class CheckCycleDirectedGraph(object):
     def __init__(self):
         self.graph = {}
         self.count = 0
 
     def print_graph(self):
         ''' for printing the contents of the graph '''
-        for i  in self.graph:
-            print(i,'->',' -> '.join([str(j) for j in self.graph[i]]))
+        for i in self.graph:
+            print(i, '->', ' -> '.join([str(j) for j in self.graph[i]]))
 
     def add_edge(self, from_vertex, to_vertex):
         ''' function to add an edge in the graph '''
@@ -114,18 +174,19 @@ def get_code(self):
         import inspect
         return inspect.getsource(CheckCycleDirected)
 
+
 class CheckCycleUndirectedGraph(object):
     def __init__(self):
         self.graph = {}
         self.count = 0
 
     def print_graph(self):
         ''' for printing the contents of the graph '''
-        for i  in self.graph:
-            print(i,'->',' -> '.join([str(j) for j in self.graph[i]]))
+        for i in self.graph:
+            print(i, '->', ' -> '.join([str(j) for j in self.graph[i]]))
 
     def add_edge(self, fromVertex, toVertex):
-        ''' for adding the edge beween two vertices '''
+        ''' for adding the edge between two vertices '''
         # check if vertex is already present,
         if fromVertex in self.graph.keys() and toVertex in self.graph.keys():
             self.graph[fromVertex].append(toVertex)
@@ -137,7 +198,7 @@ def add_edge(self, fromVertex, toVertex):
 
     def check_cycle(self):
         ''' This function will return True if graph is cyclic else return False '''
-        visited = [False] * len(self.graph)            # Marking all vertices as not visited
+        visited = [False] * len(self.graph)  # Marking all vertices as not visited
         for vertex in range(len(self.graph)):
             # Call the recursive function only if not visited
             if visited[vertex] == False:

diff --git a/pygorithm/data_structures/heap.py b/pygorithm/data_structures/heap.py
@@ -4,6 +4,7 @@
 
 from pygorithm.data_structures import queue
 
+
 # min-heap implementation as priority queue
 class Heap(queue.Queue):
     def parent_idx(self, idx):
@@ -17,7 +18,7 @@ def right_child_idx(self, idx):
 
     def insert(self, data):
         super().enqueue(data)
-        if self.rear >= 1: # heap may need to be fixed
+        if self.rear >= 1:  # heap may need to be fixed
             self.heapify_up()
 
     def heapify_up(self):
@@ -31,7 +32,7 @@ def heapify_up(self):
         Best Case:   O(1), item is inserted at correct position, no swaps needed
         Worst Case:  O(logn), item needs to be swapped throughout all levels of tree
         '''
-        child  = self.rear
+        child = self.rear
         parent = self.parent_idx(child)
         while self.queue[child] < self.queue[self.parent_idx(child)]:
             # Swap (sift up) and update child:parent relation
@@ -42,23 +43,23 @@ def heapify_up(self):
     def pop(self):
         ''' Removes the lowest value element (highest priority, at root) from the heap '''
         min = super().dequeue()
-        if self.rear >= 1: # heap may need to be fixed
+        if self.rear >= 1:  # heap may need to be fixed
             self.heapify_down()
         return min
 
     def favorite(self, parent):
         ''' Determines which child has the highest priority by 3 cases '''
-        left  = self.left_child_idx(parent)
+        left = self.left_child_idx(parent)
         right = self.right_child_idx(parent)
 
-        if left <= self.rear and right <= self.rear: # case 1: both nodes exist
+        if left <= self.rear and right <= self.rear:  # case 1: both nodes exist
             if self.queue[left] <= self.queue[right]:
                 return left
             else:
                 return right
-        elif left <= self.rear: # case 2: only left exists
+        elif left <= self.rear:  # case 2: only left exists
             return left
-        else: # case 3: no children (if left doesn't exist, neither can the right)
+        else:  # case 3: no children (if left doesn't exist, neither can the right)
             return None
 
     def heapify_down(self):
@@ -71,8 +72,8 @@ def heapify_down(self):
         Best Case:   O(1), item is inserted at correct position, no swaps needed
         Worst Case:  O(logn), item needs to be swapped throughout all levels of tree
         '''
-        cur = ROOT = 0 # start at the root
-        fav = self.favorite(cur) # determine favorite child
+        cur = ROOT = 0  # start at the root
+        fav = self.favorite(cur)  # determine favorite child
         while self.queue[fav] is not None:
             if self.queue[cur] > self.queue[fav]:
                 # Swap (sift down) and update parent:favorite relation

diff --git a/pygorithm/data_structures/linked_list.py b/pygorithm/data_structures/linked_list.py
@@ -5,16 +5,18 @@
 
 class Node(object):
     # Each node has its data and a pointer that points to next node in the Linked List
-    def __init__(self, data, next = None):
+    def __init__(self, data, next=None):
         ''' constructor '''
         self.data = data
         self.next = next
 
     # easily retrieve the source code of the Node class
-    def get_code(self):
+    @classmethod
+    def get_code(cls):
         ''' return the code for the current class '''
         import inspect
-        return inspect.getsource(Node)
+        return inspect.getsource(cls)
+
 
 class SinglyLinkedList(object):
     # Defining the head of the linked list
@@ -34,7 +36,7 @@ def get_data(self):
         ''' prints the elements in the linked list '''
         temp = self.head
         List = []
-        while(temp):
+        while (temp):
             # print(temp.data, end=' ')
             List.append(temp.data)
             temp = temp.next
@@ -43,7 +45,7 @@ def get_data(self):
 
     def insert_at_start(self, data):
         ''' insert an item at the beginning of the linked list '''
-        if self.head == None:
+        if self.head is None:
             newNode = Node(data)
             self.head = newNode
         else:
@@ -62,25 +64,25 @@ def insert_at_end(self, data):
         ''' insert an item at the end of the linked list '''
         newNode = Node(data)
         temp = self.head
-        while(temp.next != None):         # get last node
+        while (temp.next != None):  # get last node
             temp = temp.next
         temp.next = newNode
 
     def delete(self, data):
         ''' to delete specified element from the linked list '''
         temp = self.head
         # if data/key is found in head node itself
-        if (temp is not None):
-            if(temp.data == data):
+        if temp is not None:
+            if temp.data == data:
                 self.head = temp.next
                 temp = None
                 return
             else:
                 #  else _search all the nodes
-                while(temp.next != None):
-                    if(temp.data == data):
+                while (temp.next != None):
+                    if (temp.data == data):
                         break
-                    prev = temp          #save current node as previous so that we can go on to next node
+                    prev = temp  # save current node as previous so that we can go on to next node
                     temp = temp.next
 
                 # node not found
@@ -91,11 +93,13 @@ def delete(self, data):
                 return
 
     # easily retrieve the source code of the SinglyLinkedList class
-    def get_code(self):
+    @staticmethod
+    def get_code():
         ''' return the code for the current class '''
         import inspect
         return inspect.getsource(SinglyLinkedList)
 
+
 class DoublyLinkedList(object):
     def __init__(self):
         ''' constructor '''
@@ -105,7 +109,7 @@ def get_data(self):
         ''' prints the elements in the linked list '''
         temp = self.head
         List = []
-        while(temp):
+        while (temp):
             # print(temp.data, end=' ')
             List.append(temp.data)
             temp = temp.next
@@ -127,27 +131,27 @@ def insert_at_end(self, data):
         ''' insert an element at the end of the linked list '''
         newNode = Node(data)
         temp = self.head
-        while(temp.next != None):
+        while (temp.next != None):
             temp = temp.next
         temp.next = newNode
         newNode.previous = temp
 
     def delete(self, data):
         ''' to delete specified element from the linked list '''
         temp = self.head
-        if(temp.next != None):
+        if (temp.next != None):
             # if head node is to be deleted
-            if(temp.data == data):
+            if (temp.data == data):
                 temp.next.previous = None
                 self.head = temp.next
                 temp.next = None
                 return
             else:
-                while(temp.next != None):
-                    if(temp.data == data):
+                while (temp.next != None):
+                    if (temp.data == data):
                         break
                     temp = temp.next
-                if(temp.next):
+                if (temp.next):
                     # if element to be deleted is in between
                     temp.previous.next = temp.next
                     temp.next.previous = temp.previous
@@ -159,11 +163,12 @@ def delete(self, data):
                     temp.previous = None
                 return
 
-        if (temp == None):
+        if temp is None:
             return
 
     # easily retrieve the source code of the DoublyLinkedList class
-    def get_code(self):
+    @staticmethod
+    def get_code():
         ''' returns the code of the current class '''
         import inspect
         return inspect.getsource(DoublyLinkedList)
diff --git a/pygorithm/data_structures/modules.py b/pygorithm/data_structures/modules.py
@@ -1,4 +1,6 @@
 import pkgutil
+
+
 def modules():
     """
     Find all functions in pygorithm.data_structures
@@ -11,4 +13,3 @@ def modules():
     modules.remove('modules')
     modules.sort()
     return modules
-