An more sophisticated flow problem example?

[root-11]

Check if this could belong in examples. @04t02

from collections import defaultdict

from graph import Graph


def calculate_max_traffic(graph, unit_quantities):
    max_traffic = defaultdict(int)

    for in_node, out_nodes in unit_quantities.items():
        for out_node, quantity in out_nodes.items():
            _, path = graph.shortest_path(in_node, out_node)

            path_steps = list(zip(path[:-1], path[1:]))

            for (first_node, second_node) in path_steps:
                max_traffic[(first_node, second_node)] += quantity

    return max_traffic


# Example usage:
graph_edges = ["in_1", "out_1", "in_2", "out_2", "in_3", "out_3", "in_4", "out_4"]
graph_edges_pairs = list(zip(graph_edges[:-1], graph_edges[1:]))
graph_edges_pairs.append((graph_edges[-1], graph_edges[0]))

g = Graph(from_list=graph_edges_pairs)
unit_quantities = {"in_1": {"out_1": 10, "out_2": 5, "out_3": 5, "out_4": 5},  # 25
                   "in_2": {"out_2": 10, "out_3": 5, "out_4": 5, "out_1": 5},  # 25
                   "in_3": {"out_3": 5, "out_4": 5, "out_1": 5, "out_2": 10},  # 25
                   "in_4": {"out_4": 5, "out_1": 5, "out_2": 10, "out_3": 5}}  # 25

max_traffic = calculate_max_traffic(g, unit_quantities)

for (first_node, second_node), quantity in max_traffic.items():
    print(f"{first_node} -> {second_node}: {quantity}")