From 06f3057ebcbaba862f4ec303aa77b62ccbe8827c Mon Sep 17 00:00:00 2001 From: Laurent Perron Date: Tue, 6 Sep 2022 14:48:40 +0200 Subject: [PATCH] finish to numpify simple max flow --- ortools/graph/python/max_flow.cc | 3 +++ ortools/graph/samples/simple_max_flow_program.py | 10 +++++----- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/ortools/graph/python/max_flow.cc b/ortools/graph/python/max_flow.cc index f9ae45d190..742325550b 100644 --- a/ortools/graph/python/max_flow.cc +++ b/ortools/graph/python/max_flow.cc @@ -32,6 +32,8 @@ PYBIND11_MODULE(max_flow, m) { pybind11::vectorize(&SimpleMaxFlow::AddArcWithCapacity)); smf.def("set_arc_capacity", &SimpleMaxFlow::SetArcCapacity, arg("arc"), arg("capacity")); + smf.def("set_arcs_capacity", + pybind11::vectorize(&SimpleMaxFlow::SetArcCapacity)); smf.def("num_nodes", &SimpleMaxFlow::NumNodes); smf.def("num_arcs", &SimpleMaxFlow::NumArcs); smf.def("tail", &SimpleMaxFlow::Tail, arg("arc")); @@ -40,6 +42,7 @@ PYBIND11_MODULE(max_flow, m) { smf.def("solve", &SimpleMaxFlow::Solve, arg("source"), arg("sink")); smf.def("optimal_flow", &SimpleMaxFlow::OptimalFlow); smf.def("flow", &SimpleMaxFlow::Flow, arg("arc")); + smf.def("flows", pybind11::vectorize(&SimpleMaxFlow::Flow)); smf.def("get_source_side_min_cut", [](SimpleMaxFlow* smf) { std::vector result; smf->GetSourceSideMinCut(&result); diff --git a/ortools/graph/samples/simple_max_flow_program.py b/ortools/graph/samples/simple_max_flow_program.py index 5c7bf57867..eedf71d2b9 100755 --- a/ortools/graph/samples/simple_max_flow_program.py +++ b/ortools/graph/samples/simple_max_flow_program.py @@ -40,7 +40,7 @@ def main(): # [START constraints] # Add arcs in bulk. # note: we could have used add_arc_with_capacity(start, end, capacity) - smf.add_arcs_with_capacity(start_nodes, end_nodes, capacities) + all_arcs = smf.add_arcs_with_capacity(start_nodes, end_nodes, capacities) # [END constraints] # [START solve] @@ -55,10 +55,10 @@ def main(): exit(1) print('Max flow:', smf.optimal_flow()) print('') - print(' Arc Flow / Capacity') - for i in range(smf.num_arcs()): - print('%1s -> %1s %3s / %3s' % - (smf.tail(i), smf.head(i), smf.flow(i), smf.capacity(i))) + print(' Arc Flow / Capacity') + solution_flows = smf.flows(all_arcs) + for arc, flow, capacity in zip(all_arcs, solution_flows, capacities): + print(f'{smf.tail(arc)} / {smf.head(arc)} {flow:3} / {capacity:3}') print('Source side min-cut:', smf.get_source_side_min_cut()) print('Sink side min-cut:', smf.get_sink_side_min_cut()) # [END print_solution]