use std::collections::BTreeMap;
use std::iter;
use std::cmp::min;
#[link(name="flow")]
extern {
fn network_simplex_mcmf_i64(num_vertices: i64, num_edges: i64,
node_supply: *const i64, edge_a: *const i64, edge_b: *const i64,
edge_capacity: *const i64, edge_cost: *const i64, edge_flow_result: *mut i64) -> i64;
}
#[derive(Clone, Copy, Eq, PartialEq)]
struct Node(usize);
struct Edge {
pub a: Node,
pub b: Node,
pub data: EdgeData,
}
struct Graph {
nodes: Vec<NodeData>,
edges: Vec<Edge>,
}
impl Graph {
pub fn add_edge(&mut self, a: Node, b: Node, data: EdgeData) -> &mut Self {
assert!(a.0 < self.nodes.len());
assert!(b.0 < self.nodes.len());
self.edges.push(Edge {a, b, data});
self
}
pub fn extract(self) -> (Vec<NodeData>, Vec<Edge>) {
(self.nodes, self.edges)
}
}
impl Graph {
pub fn new_default(num_vertices: usize) -> Self {
let nodes = vec![Default::default(); num_vertices];
Graph {nodes, edges: Vec::new()}
}
}
#[derive(Clone, Copy, Default)]
struct NodeData {
supply: i64,
}
#[derive(Clone, Copy)]
struct EdgeData {
cost: i64,
capacity: i64,
flow: i64,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Cost(pub i32);
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Capacity(pub i32);
impl EdgeData {
pub fn new(cost: Cost, capacity: Capacity) -> Self {
let cost = cost.0 as i64;
let capacity = capacity.0 as i64;
assert!(capacity >= 0);
EdgeData {cost, capacity, flow: Default::default()}
}
}
impl Graph {
pub fn increase_supply(&mut self, node: Node, amount: i64) {
self.delta_supply(node, amount);
}
pub fn decrease_supply(&mut self, node: Node, amount: i64) {
self.delta_supply(node, -amount);
}
pub fn delta_supply(&mut self, node: Node, amount: i64) {
self.nodes[node.0].supply += amount;
}
pub fn mcmf(&mut self) -> i64 {
let num_vertices = self.nodes.len() as i64;
let num_edges = self.edges.len() as i64;
let node_supply: Vec<_> = self.nodes.iter().map(|x| clamp_to_i32(x.supply)).collect();
let edge_a: Vec<_> = self.edges.iter().map(|x| x.a.0 as i64).collect();
let edge_b: Vec<_> = self.edges.iter().map(|x| x.b.0 as i64).collect();
let edge_capacity: Vec<_> = self.edges.iter().map(|x| x.data.capacity).collect();
let edge_cost: Vec<_> = self.edges.iter().map(|x| x.data.cost).collect();
let mut edge_flow_result = vec![0; self.edges.len()];
let result;
unsafe {
result = network_simplex_mcmf_i64(num_vertices, num_edges,
node_supply.as_ptr(), edge_a.as_ptr(), edge_b.as_ptr(),
edge_capacity.as_ptr(), edge_cost.as_ptr(), edge_flow_result.as_mut_ptr());
}
for (edge, &flow) in self.edges.iter_mut().zip(edge_flow_result.iter()) {
edge.data.flow = flow;
}
result
}
}
fn clamp_to_i32(x: i64) -> i64 {
let limit = std::i32::MAX as i64;
let x = std::cmp::min(x, limit);
let x = std::cmp::max(x, -limit);
x
}
#[derive(Copy, Clone, Hash, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Vertex<T: Clone + Ord> {
Source, Sink, Node(T)
}
impl<T> Vertex<T> where T: Clone + Ord {
pub fn as_option(self) -> Option<T> {
match self {
Vertex::Source => None,
Vertex::Sink => None,
Vertex::Node(x) => Some(x),
}
}
}
impl<T> From<T> for Vertex<T> where T: Clone + Ord {
fn from(x: T) -> Vertex<T> {
Vertex::Node(x)
}
}
#[derive(Clone)]
pub struct Flow<T: Clone + Ord> {
pub a: Vertex<T>,
pub b: Vertex<T>,
pub amount: u32,
pub cost: i32
}
pub struct Path<T: Clone + Ord> {
pub flows: Vec<Flow<T>>
}
impl<T> Path<T> where T: Clone + Ord {
pub fn vertices(&self) -> Vec<&Vertex<T>> {
iter::once(&self.flows[0].a)
.chain(self.flows.iter().map(|x| &x.b))
.collect()
}
pub fn edges(&self) -> Vec<&Flow<T>> {
self.flows.iter().collect()
}
pub fn cost(&self) -> i32 {
self.flows.iter()
.map(|flow| flow.amount as i32 * flow.cost)
.sum()
}
pub fn amount(&self) -> u32 {
self.flows[0].amount
}
pub fn len(&self) -> usize {
self.flows.len()
}
}
#[derive(Clone)]
pub struct GraphBuilder<T: Clone + Ord> {
pub edge_list: Vec<(Vertex<T>, Vertex<T>, Capacity, Cost)>
}
impl<T> GraphBuilder<T> where T: Clone + Ord {
pub fn new() -> Self {
GraphBuilder {edge_list: Vec::new()}
}
pub fn add_edge<A: Into<Vertex<T>>, B: Into<Vertex<T>>>(&mut self, a: A, b: B, capacity: Capacity, cost: Cost) -> &mut Self {
if capacity.0 < 0 {
panic!("capacity cannot be negative (capacity was {})", capacity.0)
}
let a = a.into();
let b = b.into();
assert!(a != b);
assert!(a != Vertex::Sink);
assert!(b != Vertex::Source);
self.edge_list.push((a, b, capacity, cost));
self
}
pub fn mcmf(&self) -> (i32, Vec<Path<T>>) {
let mut next_id = 0;
let source = Vertex::Source.clone();
let sink = Vertex::Sink.clone();
let mut index_mapper = BTreeMap::new();
for vertex in self.edge_list.iter()
.flat_map(move |&(ref a, ref b, _, _)| iter::once(a).chain(iter::once(b)))
.chain(iter::once(&source))
.chain(iter::once(&sink)) {
if !index_mapper.contains_key(&vertex) {
index_mapper.insert(vertex, next_id);
next_id += 1;
}
}
let num_vertices = next_id;
let mut g = Graph::new_default(num_vertices);
for &(ref a, ref b, cap, cost) in &self.edge_list {
let node_a = Node(*index_mapper.get(&a).unwrap());
let node_b = Node(*index_mapper.get(&b).unwrap());
if *a == Vertex::Source || *b == Vertex::Sink {
g.increase_supply(Node(*index_mapper.get(&Vertex::Source).unwrap()), cap.0 as i64);
g.decrease_supply(Node(*index_mapper.get(&Vertex::Sink).unwrap()), cap.0 as i64);
}
g.add_edge(node_a, node_b, EdgeData::new(cost, cap));
}
let total_amount = g.mcmf() as i32;
let (_, edges) = g.extract();
let inverse_mapping: BTreeMap<_, _> =
index_mapper.into_iter().map(|(a, b)| (b, a)).collect();
let flows = edges.into_iter().map(|x| {
let a = (**inverse_mapping.get(&x.a.0).unwrap()).clone();
let b = (**inverse_mapping.get(&x.b.0).unwrap()).clone();
let amount = x.data.flow as u32;
let cost = x.data.cost as i32;
Flow {a, b, amount, cost}
})
.filter(|x| x.amount != 0)
.collect();
let mut paths = GraphBuilder::path_decomposition(flows);
paths.sort_by_key(|path| path.len());
(total_amount, paths)
}
fn path_decomposition(flows: Vec<Flow<T>>) -> Vec<Path<T>> {
let mut adj: BTreeMap<Vertex<T>, Vec<Flow<T>>> = flows.iter()
.map(|x| (x.a.clone(), Vec::new()))
.collect();
for x in flows {
adj.get_mut(&x.a).unwrap().push(x);
}
fn decompose<T: Clone + Ord>(adj: &mut BTreeMap<Vertex<T>, Vec<Flow<T>>>, v: &Vertex<T>, parent_amount: u32) -> (u32, Vec<Flow<T>>) {
if *v == Vertex::Sink {
(std::u32::MAX, Vec::new())
} else if adj.get(&v).into_iter().all(|x| x.is_empty()) {
(0, Vec::new())
} else {
let flow = adj.get_mut(&v).unwrap().pop().unwrap();
let amount = min(parent_amount, flow.amount);
let (child_amount, child_path) = decompose(adj, &flow.b, amount);
let amount = min(amount, child_amount);
let mut path = child_path;
if amount < flow.amount {
adj.get_mut(&v).unwrap().push(Flow {amount: flow.amount - amount, ..flow.clone()});
}
path.push(Flow {amount, ..flow});
(amount, path)
}
}
let mut result = Vec::new();
loop {
let (flow, path) = decompose(&mut adj, &Vertex::Source, std::u32::MAX);
if flow == 0 {
break;
} else {
result.push(path.into_iter().rev().collect());
}
}
result.into_iter().map(|x| Path {flows: x}).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[allow(non_snake_case)]
fn mcmf() {
let mut G = Graph::new_default(4);
G.increase_supply(Node(0), 20);
G.decrease_supply(Node(3), 20);
G.add_edge(Node(0), Node(1), EdgeData::new(Cost(100), Capacity(10)));
G.add_edge(Node(0), Node(2), EdgeData::new(Cost(300), Capacity(20)));
G.add_edge(Node(1), Node(2), EdgeData::new(Cost(50), Capacity(5)));
G.add_edge(Node(1), Node(3), EdgeData::new(Cost(200), Capacity(10)));
G.add_edge(Node(2), Node(3), EdgeData::new(Cost(100), Capacity(20)));
let cost = G.mcmf();
let (_, edges) = G.extract();
let flow: Vec<_> = edges.iter().map(|x| x.data.flow).collect();
assert_eq!(cost, 6750);
assert_eq!(flow, vec![10, 10, 5, 5, 15]);
}
#[test]
fn large_number() {
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
enum OnlyNode {
Only
}
for i in 0..48 {
let x = i * 1000;
println!("x={}", x);
let (total, _) = GraphBuilder::new()
.add_edge(Vertex::Source, OnlyNode::Only, Capacity(x), Cost(x))
.add_edge(Vertex::Source, OnlyNode::Only, Capacity(x), Cost(x))
.add_edge(OnlyNode::Only, Vertex::Sink, Capacity(x), Cost(0))
.mcmf();
assert_eq!(total, (x as i64 * x as i64) as i32);
}
}
#[test]
fn empty_graph() {
let (cost, paths) = GraphBuilder::<i32>::new().mcmf();
assert_eq!(cost, 0);
assert!(paths.is_empty())
}
#[test]
fn large_capacities() {
let max = 1 << 30;
assert_eq!(max, 1073741824);
let (cost, paths) = GraphBuilder::new()
.add_edge(
Vertex::Source,
"A",
Capacity(max),
Cost(0),
).add_edge(
"A",
Vertex::Sink,
Capacity(max),
Cost(0),
).mcmf();
assert_eq!(cost, 0);
assert_eq!(paths.len(), 1);
}
#[test]
#[should_panic]
fn negative_capacity_panics() {
GraphBuilder::new().add_edge("a", "b", Capacity(-1), Cost(0));
}
}