Add doctests to networking_flow/minimum_cut.py

networking_flow/minimum_cut.py

@@ -1,12 +1,21 @@
 # Minimum cut on Ford_Fulkerson algorithm.
-
+
+test_graph = [
+    [0, 16, 13, 0, 0, 0],
+    [0, 0, 10, 12, 0, 0],
+    [0, 4, 0, 0, 14, 0],
+    [0, 0, 9, 0, 0, 20],
+    [0, 0, 0, 7, 0, 4],
+    [0, 0, 0, 0, 0, 0],
+]
+
+
 def BFS(graph, s, t, parent):
     # Return True if there is node that has not iterated.
-    visited = [False]*len(graph)
-    queue=[]
-    queue.append(s)
+    visited = [False] * len(graph)
+    queue = [s]
     visited[s] = True
-    
+
     while queue:
         u = queue.pop(0)
         for ind in range(len(graph[u])):
@@ -16,26 +25,30 @@ def BFS(graph, s, t, parent):
                 parent[ind] = u
 
     return True if visited[t] else False
-
+
+
 def mincut(graph, source, sink):
-    # This array is filled by BFS and to store path
-    parent = [-1]*(len(graph))
-    max_flow = 0 
+    """This array is filled by BFS and to store path
+    >>> mincut(test_graph, source=0, sink=5)
+    [(1, 3), (4, 3), (4, 5)]
+    """
+    parent = [-1] * (len(graph))
+    max_flow = 0
     res = []
-    temp = [i[:] for i in graph]   # Record orignial cut, copy.
-    while BFS(graph, source, sink, parent) :
+    temp = [i[:] for i in graph]  # Record orignial cut, copy.
+    while BFS(graph, source, sink, parent):
         path_flow = float("Inf")
         s = sink
 
-        while(s !=  source):
+        while s != source:
             # Find the minimum value in select path
-            path_flow = min (path_flow, graph[parent[s]][s])
+            path_flow = min(path_flow, graph[parent[s]][s])
             s = parent[s]
 
-        max_flow +=  path_flow
+        max_flow += path_flow
         v = sink
-        
-        while(v !=  source):
+
+        while v != source:
             u = parent[v]
             graph[u][v] -= path_flow
             graph[v][u] += path_flow
@@ -44,16 +57,10 @@ def mincut(graph, source, sink):
     for i in range(len(graph)):
         for j in range(len(graph[0])):
             if graph[i][j] == 0 and temp[i][j] > 0:
-                res.append((i,j))
+                res.append((i, j))
 
     return res
 
-graph = [[0, 16, 13, 0, 0, 0],
-         [0, 0, 10 ,12, 0, 0],
-         [0, 4, 0, 0, 14, 0],
-         [0, 0, 9, 0, 0, 20],
-         [0, 0, 0, 7, 0, 4],
-         [0, 0, 0, 0, 0, 0]]
 
-source, sink = 0, 5
-print(mincut(graph, source, sink))
+if __name__ == "__main__":
+    print(mincut(test_graph, source=0, sink=5))