use crate::traits::IndexDigraph;
use std::cmp::min;
use std::collections::VecDeque;
use crate::num::traits::NumAssign;
pub struct Dinic<'a, G, F>
where
G: IndexDigraph,
{
g: &'a G,
nodes: Vec<NodeInfo>,
neighs: Vec<Vec<(usize, usize)>>,
edges: Vec<EdgeInfo<F>>,
queue: VecDeque<usize>,
value: F,
}
#[derive(Clone)]
struct NodeInfo {
dist: usize,
first_lvl: (usize, usize),
}
#[derive(Clone)]
struct EdgeInfo<F> {
flow: F,
next_lvl: (usize, usize),
}
impl<'a, G, F> Dinic<'a, G, F>
where
G: IndexDigraph,
F: NumAssign + Ord + Copy,
{
pub fn new(g: &'a G) -> Self {
Dinic {
g,
nodes: vec![
NodeInfo {
dist: 0,
first_lvl: (usize::max_value(), usize::max_value()),
};
g.num_nodes()
],
neighs: g
.nodes()
.map(|u| {
g.outedges(u)
.map(|(e, v)| (g.edge_id(e) << 1, g.node_id(v)))
.chain(g.inedges(u).map(|(e, v)| ((g.edge_id(e) << 1) | 1, g.node_id(v))))
.collect()
})
.collect(),
edges: vec![
EdgeInfo {
flow: F::zero(),
next_lvl: (usize::max_value(), usize::max_value()),
};
g.num_edges() * 2
],
queue: VecDeque::with_capacity(g.num_nodes()),
value: F::zero(),
}
}
pub fn as_graph(&self) -> &'a G {
self.g
}
pub fn value(&self) -> F {
self.value
}
pub fn flow(&self, e: G::Edge<'_>) -> F {
self.edges[self.g.edge_id(e) << 1].flow
}
pub fn flow_iter<'b>(&'b self) -> impl Iterator<Item = (G::Edge<'a>, F)> + 'b {
self.g
.edges()
.enumerate()
.map(move |(i, e)| (e, self.edges[i << 1].flow))
}
pub fn solve<Us>(&mut self, src: G::Node<'_>, snk: G::Node<'_>, upper: Us)
where
Us: Fn(G::Edge<'a>) -> F,
{
let src = self.g.node_id(src);
let snk = self.g.node_id(snk);
assert_ne!(src, snk, "Source and sink node must not be equal");
for (e, einfo) in self.edges.iter_mut().enumerate() {
einfo.flow = if (e & 1) == 0 {
F::zero()
} else {
upper(self.g.id2edge(e >> 1))
};
}
self.value = F::zero();
while self.search(src, snk) {
let v = self.augment(src, snk, None);
self.value += v;
}
}
pub fn mincut(&self) -> Vec<G::Node<'a>> {
let n = self.g.num_nodes();
self.g
.nodes()
.filter(|&u| self.nodes[self.g.node_id(u)].dist < n)
.collect()
}
fn search(&mut self, src: usize, snk: usize) -> bool {
let n = self.g.num_nodes();
for node in &mut self.nodes {
node.dist = n;
node.first_lvl = (usize::max_value(), usize::max_value());
}
self.nodes[src].dist = 0;
self.queue.clear();
self.queue.push_back(src);
let mut snk_d = n;
while let Some(u) = self.queue.pop_front() {
let d = self.nodes[u].dist;
if d >= snk_d {
return true;
}
for &(e, v) in &self.neighs[u] {
if self.edges[e ^ 1].flow > F::zero() {
if self.nodes[v].dist == n {
self.nodes[v].dist = d + 1;
self.queue.push_back(v);
if v == snk {
snk_d = d + 1
}
} else if self.nodes[v].dist != d + 1 {
continue;
}
self.edges[e].next_lvl = self.nodes[u].first_lvl;
self.nodes[u].first_lvl = (e, v);
}
}
}
false
}
fn augment(&mut self, src: usize, snk: usize, target_flow: Option<F>) -> F {
if src == snk {
return target_flow.expect("Unbounded flow (are source and sink the same?)");
}
let mut df = F::zero();
loop {
let (e, v) = self.nodes[src].first_lvl;
if e == usize::max_value() {
break;
}
let f = e ^ 1;
let rem_cap = match target_flow {
Some(target_flow) => min(self.edges[f].flow, target_flow - df),
None => self.edges[f].flow,
};
if rem_cap > F::zero() {
let cf = self.augment(v, snk, Some(rem_cap));
self.edges[e].flow += cf;
self.edges[f].flow -= cf;
df += cf;
if target_flow.map(|t| df == t).unwrap_or(false) {
break;
}
}
self.nodes[src].first_lvl = self.edges[e].next_lvl;
}
if df.is_zero() {
self.nodes[src].first_lvl = (usize::max_value(), usize::max_value());
}
df
}
}
pub fn dinic<'a, G, F, Us>(
g: &'a G,
src: G::Node<'_>,
snk: G::Node<'_>,
upper: Us,
) -> (F, Vec<(G::Edge<'a>, F)>, Vec<G::Node<'a>>)
where
G: IndexDigraph,
F: 'a + NumAssign + Ord + Copy,
Us: Fn(G::Edge<'_>) -> F,
{
let mut maxflow = Dinic::new(g);
maxflow.solve(src, snk, upper);
(maxflow.value(), maxflow.flow_iter().collect(), maxflow.mincut())
}