Dinics maximum flow in directed graph algorithm

dinics-algorithm.py

# Dinics Algorithm
# - Find level graph (DAG)
# - Find blocking flows in level graph
# - Continue as long as no more flow can be added
# - O(EV^2) time O(EV) for each phase in O(V) level graphs

INF = 10**9

class Edge:
  def __init__(self, u, v, cap):
    self.u = u
    self.v = v
    self.cap = cap
    self.flow = 0
    self.residual = None
    
  def augment(self, flow):
    self.flow += flow
    self.residual.flow -= flow
    
  def get_cap(self):
    return self.cap - self.flow
    
graph = dict() # node -> list[Edge]

def add_edge(u, v, cap):
  edge = Edge(u, v, cap)
  reverse = Edge(v, u, 0)
  edge.residual = reverse
  reverse.residual = edge
  graph[u].append(edge)
  graph[v].append(reverse)
  
def max_flow(graph, n, source, target):
  levels = [-1] * n
  start = [0] * n
  
  max_f = 0
  while bfs(levels, source, target):
    f = dfs(levels, start, source, target)
    if f == 0: break
    max_f += f
    levels = [-1] * n
    
  return max_f

def bfs(levels, source, target):
  levels[source] = 0
  q = deque([source])
  while q:
    u = q.popleft()
    for edge in graph[u]:
      if edge.get_cap() > 0 and levels[edge.v] == -1:
        levels[edge.v] = levels[u] + 1
        q.append(edge.v)
      
  return levels[target] != -1

def dfs(levels, start, source, target, flow=INF):
  if flow == 0: return 0
  if source == target: return flow
  
  for i in range(start[source], len(graph[source])):
    edge = graph[source][i]
    if edge.get_cap() > 0 and levels[edge.v] == levels[source] + 1:
      f = dfs(levels, start, edge.v, target, min(flow, edge.get_cap()))
      if f > 0:
        edge.augment(f)
        return f
    start[source] += 1
    
  return 0