Dinitz correction

networkx/algorithms/flow/dinitz_alg.py

@@ -173,35 +173,64 @@ def dinitz_impl(G, s, t, capacity, residual, cutoff):
 
     def breath_first_search():
         parents = {}
-        queue = deque([s])
+        vertex_dist = {s: 0}
+        queue = deque([(s, 0)])
+        # Record all the potential edges of shortest augmenting paths
         while queue:
             if t in parents:
                 break
-            u = queue.popleft()
-            for v in R_succ[u]:
-                attr = R_succ[u][v]
-                if v not in parents and attr["capacity"] - attr["flow"] > 0:
-                    parents[v] = u
-                    queue.append(v)
+            u, dist = queue.popleft()
+            for v, attr in R_succ[u].items():
+                if attr["capacity"] - attr["flow"] > 0:
+                    if v in parents:
+                        if vertex_dist[v] == dist + 1:
+                            parents[v].append(u)
+                    else:
+                        parents[v] = deque([u])
+                        vertex_dist[v] = dist + 1
+                        queue.append((v, dist + 1))
         return parents
 
     def depth_first_search(parents):
+        # DFS to find all the shortest augmenting paths
         """Build a path using DFS starting from the sink"""
-        path = []
+        total_flow = 0
         u = t
-        flow = INF
-        while u != s:
-            path.append(u)
-            v = parents[u]
-            flow = min(flow, R_pred[u][v]["capacity"] - R_pred[u][v]["flow"])
+        # path also functions as a stack
+        path = [u]
+        # The loop ends with no augmenting path left in the layered graph
+        while True:
+            if len(parents[u]) > 0:
+                v = parents[u][0]
+                path.append(v)
+            else:
+                path.pop()
+                if len(path) == 0:
+                    break
+                v = path[-1]
+                parents[v].popleft()
+            # Augment the flow along the path found
+            if v == s:
+                flow = INF
+                for u, v in pairwise(path):
+                    flow = min(flow, R_pred[u][v]["capacity"] - R_pred[u][v]["flow"])
+                u_index = len(path) - 2
+                while u_index >= 0:
+                    u = path[u_index]
+                    v_index = u_index + 1
+                    v = path[v_index]
+                    R_pred[u][v]["flow"] += flow
+                    R_pred[v][u]["flow"] -= flow
+                    # Find the proper node to continue the search
+                    if R_pred[u][v]["capacity"] - R_pred[u][v]["flow"] == 0:
+                        parents[u].popleft()
+                        while path[-1] != u:
+                            path.pop()
+                    u_index -= 1
+                total_flow += flow
+                v = path[-1]
             u = v
-        path.append(s)
-        # Augment the flow along the path found
-        if flow > 0:
-            for u, v in pairwise(path):
-                R_pred[u][v]["flow"] += flow
-                R_pred[v][u]["flow"] -= flow
-        return flow
+        return total_flow
 
     flow_value = 0
     while flow_value < cutoff: