AAM: Algorithms for Applied Mathematics

def all_shortest_paths(vertex_set, warrows, source):
    # Note: float('inf') represents Infinity.
    unvisited_distance = [(0,source)]
    unvisited_distance += [(float('inf'), v) for v in vertex_set if v!=source]
    distance = {v:d for d,v in unvisited_distance}
    predecessor = {v:[] for v in vertex_set}
    visited = [ ]
    while len(visited) < len(vertex_set):
        # find nearest unvisited vertex
        unvisited_distance.sort()
        du, u = unvisited_distance.pop(0)
        for v, cap_uv in warrows[u].items():
            if v in visited: continue
            elif (du + 1 <= distance[v]) and (cap_uv > 0):
                unvisited_distance.remove( (distance[v], v) )
                distance[v] = du + 1
                if u not in predecessor[v]:
                    predecessor[v].append( u )
                unvisited_distance.append( (distance[v], v) )
        visited.append( u )
    return predecessor

def build_back_arrows(vertex_set, warrows, predecessors):
    back_arrows = {v:{} for v in vertex_set}
    for u,pu in predecessors.items():
        for v in pu:
            try:
                back_arrows[u][v] = warrows[v][u]
            except KeyError:
                back_arrows[u] = {v:warrows[v][u]}
    return back_arrows

def pathfinding(warrows:dict, source, target):
    augpath = [source]
    augflow = float('inf')
    while augpath[-1] != target:
        neighbors = list(warrows[augpath[-1]].items())
        next_vert, next_cap = neighbors[0]
        if next_cap < augflow:
            augflow = next_cap
        augpath.append(next_vert)
    return augpath, augflow

def predecessors(warrows:dict, vertex):
    return [w for w,nbrs in warrows.items() if vertex in nbrs.keys() and warrows[w][vertex] > 0 ]

def successors(warrows:dict, vertex):
    return [w for w,c in warrows[vertex].items() if c>0]

def clean_to_source(layered_cap, vertex, source):
    if vertex!=source and len(successors(layered_cap, vertex)) == 0:
        for w in predecessors(layered_cap, vertex):
            layered_cap[w][vertex] = 0
            if len(successors(layered_cap, w))==0:
                layered_cap = clean_to_source(layered_cap, w, source)
    return layered_cap

def clean_to_target(layered_cap, vertex, target):
    if vertex!=target and len(predecessors(layered_cap, vertex))==0:
        for w in successors(layered_cap, vertex):
            layered_cap[vertex][w] = 0
            if predecessors(layered_cap, w) in ([vertex], []):
                layered_cap = clean_to_target(layered_cap, w, target)
    return layered_cap

def maintain_layers(layered_cap, aug_path, aug_flow, source, target):
    pairs = zip(aug_path[:-1], aug_path[1:])
    for u,v in pairs:
        layered_cap[u][v] -= aug_flow
        if layered_cap[u][v] == 0:
            layered_cap = clean_to_source(layered_cap, u, source)
            layered_cap = clean_to_target(layered_cap, v, target)
    pruners = []
    for u,nbrs in layered_cap.copy().items():
        for v,c in nbrs.items():
            if c==0:
                pruners.append( (u,v) )
    for u,v in pruners:
        del layered_cap[u][v]
    return layered_cap

def dinitz(vertex_set, capacity, source, target, debug=False):
    cap_edges = [(u,v,c) for u,out_arrows in capacity.items() for v,c in out_arrows.items()]
    # Initialize Zero Flow
    total_flow = 0
    flow = {u:{v:0 for v in out_arrows.keys()} for u,out_arrows in capacity.items()}
    for u,v,c in cap_edges:
        try:
            flow[v][u] = 0
        except:
            flow[v] = {u:0}
    while True:
        # Create residual capacity
        res_capacity = {u:{v:c-flow[u][v] for v,c in out_arrows.items()} for u,out_arrows in capacity.items()}
        for u,v,c in cap_edges:
            try:
                res_capacity[v][u] = flow[u][v]
            except KeyError:
                res_capacity[v] = {u:flow[u][v]}

        res_capacity = {u:{v:c for v,c in out_arrows.items() if c>0} for u,out_arrows in res_capacity.items()}
        if debug: print(f"Residual Capacities computed:\n{res_capacity}")
        # Create layered graph
        pred = all_shortest_paths(vertex_set, res_capacity, source)
        back_cap = build_back_arrows(vertex_set, res_capacity, pred)
        pred = all_shortest_paths(vertex_set, back_cap, target)
        layered_capacity = build_back_arrows(vertex_set, back_cap, pred)
        if debug: print(f"Layered Graph computed:\n{layered_capacity}")
        if layered_capacity[source] == {}:
            print(f"Maximum flow of {total_flow} achieved.")
            break
        try:
            # Repeatedly find augmenting paths
            while True:
                ap, af = pathfinding(layered_capacity, source, target)
                if debug: print(f"Augmenting path {ap} of capacity {af} found.")
                total_flow += af
                for u,v in zip(ap[:-1],ap[1:]):
                    flow[u][v] += af
                layered_capacity = maintain_layers(layered_capacity, ap, af, source, target)
                if debug: print(f"After maintenance, layered graph is:\n{layered_capacity}")
        except:
            if debug: print(f"No more augmenting paths.")
            continue
    # prune empty flows
    flow = {u:{v:f for v,f in flow_arrows.items() if f>0} for u,flow_arrows in flow.items()}
    return flow