use crate::traits::IndexDigraph;
use std::cmp::min;
use std::collections::VecDeque;
use crate::num::traits::NumAssign;
pub struct PushRelabel<'a, G, F>
where
G: 'a + IndexDigraph,
{
g: &'a G,
nodes: Vec<NodeInfo<F>>,
edges: Vec<Vec<(usize, usize)>>,
flow: Vec<F>,
buckets: Vec<Bucket>,
queue: VecDeque<usize>,
largest_act: usize,
value: F,
pub cnt_relabel: usize,
pub use_global_relabelling: bool,
}
#[derive(Clone)]
struct NodeInfo<Flow> {
height: usize,
excess: Flow,
next_act: usize,
iter: usize,
}
impl<Flow> NodeInfo<Flow>
where
Flow: NumAssign,
{
fn reset(&mut self) {
self.height = 0;
self.excess = Flow::zero();
self.next_act = usize::max_value();
self.iter = 0;
}
}
#[derive(Clone)]
struct Bucket {
first_act: usize,
num_inact: usize,
}
impl Bucket {
fn is_empty(&self) -> bool {
self.first_act == usize::max_value() && self.num_inact == 0
}
}
impl<'a, G, F> PushRelabel<'a, G, F>
where
G: IndexDigraph,
F: NumAssign + Ord + Copy,
{
pub fn new(g: &'a G) -> Self {
let n = g.num_nodes();
PushRelabel {
g,
nodes: vec![
NodeInfo {
height: 0,
excess: F::zero(),
next_act: usize::max_value(),
iter: 0,
};
n
],
edges: 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(),
flow: vec![F::zero(); g.num_edges() * 2],
buckets: vec![
Bucket {
first_act: usize::max_value(),
num_inact: 0,
};
n * 2
],
queue: VecDeque::with_capacity(n),
largest_act: 0,
value: F::zero(),
cnt_relabel: 0,
use_global_relabelling: true,
}
}
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.flow[self.g.edge_id(e) << 1]
}
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.flow[i << 1]))
}
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");
let n = self.g.num_nodes();
self.cnt_relabel = 0;
self.init_preflow(src, upper);
self.update_heights(src, snk, false);
let mut lvl_relabel = if self.use_global_relabelling {
n
} else {
usize::max_value()
};
while self.largest_act != 0 && self.largest_act != n {
let l = self.largest_act;
let u = self.buckets[l].first_act;
self.buckets[l].first_act = self.nodes[u].next_act;
self.discharge(u);
if self.cnt_relabel >= lvl_relabel && self.largest_act < n {
self.update_heights(src, snk, false);
lvl_relabel += n;
}
if self.largest_act == 0 {
self.update_heights(src, snk, true);
}
}
self.value = self.nodes[snk].excess;
}
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)].height >= n)
.collect()
}
fn init_preflow<Us>(&mut self, src: usize, upper: Us)
where
Us: Fn(G::Edge<'a>) -> F,
{
for node in &mut self.nodes {
node.reset();
}
for (e, flw) in self.flow.iter_mut().enumerate() {
if e & 1 == 0 {
*flw = F::zero();
} else {
*flw = upper(self.g.id2edge(e >> 1));
}
}
for &(e, v) in &self.edges[src] {
let f = e ^ 1;
let ub = self.flow[f];
self.flow[e] = ub;
self.flow[f] = F::zero();
self.nodes[v].excess += ub;
}
self.nodes[src].height = self.g.num_nodes();
}
fn update_heights(&mut self, src: usize, snk: usize, from_src: bool) {
let n = self.g.num_nodes();
for node in &mut self.nodes {
node.height = n + n;
node.iter = 0;
}
self.nodes[snk].height = 0;
self.nodes[src].height = n;
for b in if from_src {
&mut self.buckets[n + 1..]
} else {
&mut self.buckets[..n]
} {
b.first_act = usize::max_value();
b.num_inact = 0;
}
self.buckets[0].num_inact = 1;
self.buckets[n].num_inact = 1;
self.queue.clear();
self.queue.push_back(snk);
self.largest_act = 0;
loop {
while let Some(v) = self.queue.pop_front() {
let h = self.nodes[v].height + 1;
for &(e, u) in &self.edges[v] {
let udata = &mut self.nodes[u];
if udata.height > h && self.flow[e] > F::zero() {
udata.height = h;
self.queue.push_back(u);
if udata.excess > F::zero() {
udata.next_act = self.buckets[h].first_act;
self.buckets[h].first_act = u;
self.largest_act = h;
debug_assert!(from_src || self.largest_act < n);
} else {
self.buckets[h].num_inact += 1;
}
}
}
}
if self.largest_act >= n || !from_src {
break;
}
self.largest_act = n;
self.queue.push_back(src);
}
}
#[allow(clippy::many_single_char_names)]
fn discharge(&mut self, u: usize) {
let n = self.g.num_nodes();
loop {
let mut h_neighbor = 2 * n;
let h_u = self.nodes[u].height;
let edges = &self.edges[u];
let first_iter = self.nodes[u].iter;
let mut cur = first_iter;
loop {
let (e, v) = edges[cur];
let f = e ^ 1;
if !self.flow[f].is_zero() {
if h_u == self.nodes[v].height + 1 {
let df = min(self.nodes[u].excess, self.flow[f]);
debug_assert_eq!(h_u, self.nodes[v].height + 1);
debug_assert!(df > F::zero());
if self.nodes[v].excess.is_zero() {
let h = self.nodes[v].height;
self.nodes[v].next_act = self.buckets[h].first_act;
self.buckets[h].first_act = v;
debug_assert!(self.buckets[h].num_inact > 0, "height:{} n:{}", h, n);
self.buckets[h].num_inact -= 1;
}
self.flow[e] += df;
self.flow[f] -= df;
self.nodes[u].excess -= df;
self.nodes[v].excess += df;
if self.nodes[u].excess.is_zero() {
self.buckets[h_u].num_inact += 1;
self.largest_act = (0..h_u + 1)
.rev()
.find(|&h| self.buckets[h].first_act != usize::max_value())
.unwrap_or(0);
self.nodes[u].iter = cur;
return;
}
} else {
h_neighbor = h_neighbor.min(self.nodes[v].height);
}
}
cur += 1;
if cur == edges.len() {
break;
}
}
self.nodes[u].iter = 0;
for &(e, v) in &edges[..first_iter] {
if !self.flow[e ^ 1].is_zero() {
h_neighbor = h_neighbor.min(self.nodes[v].height);
}
}
if !self.relabel(u, h_neighbor + 1) {
break;
}
}
}
fn relabel(&mut self, u: usize, h_new: usize) -> bool {
debug_assert!(self.nodes[u].excess > F::zero());
self.cnt_relabel += 1;
let n = self.g.num_nodes();
let h_old = self.nodes[u].height;
debug_assert_eq!(
h_new,
self.edges[u]
.iter()
.filter_map(|&(e, v)| {
if self.flow[e ^ 1] > F::zero() {
Some(self.nodes[v].height)
} else {
None
}
})
.min()
.unwrap()
+ 1
);
debug_assert!(h_new > h_old);
if h_old < n {
let mut h_new = h_new;
if self.buckets[h_old].is_empty() {
for h in h_old + 1..n {
if self.buckets[h].is_empty() {
break;
}
self.buckets[h].first_act = usize::max_value();
self.buckets[h].num_inact = 0;
}
for node in &mut self.nodes {
if h_old < node.height {
node.height = n + 1;
}
}
h_new = n + 1;
}
if h_new >= n {
debug_assert_eq!(self.largest_act, h_old);
self.nodes[u].height = h_new;
self.largest_act = (0..h_old + 1)
.rev()
.find(|&h| self.buckets[h].first_act != usize::max_value())
.unwrap_or(0);
return false;
}
}
self.nodes[u].height = h_new;
self.largest_act = h_new;
true
}
}
#[cfg(test)]
mod tests {
use crate::maxflow::pushrelabel;
use crate::traits::*;
use crate::{Buildable, Builder, Net};
#[test]
fn test_pushrelabel() {
let mut g = Net::new_builder();
let mut upper = vec![];
let s = g.add_node();
let t = g.add_node();
let v1 = g.add_node();
let v2 = g.add_node();
let v3 = g.add_node();
let v4 = g.add_node();
g.add_edge(s, v1);
upper.push(15);
g.add_edge(s, v3);
upper.push(10);
g.add_edge(v1, v2);
upper.push(6);
g.add_edge(v1, v3);
upper.push(7);
g.add_edge(v2, t);
upper.push(5);
g.add_edge(v2, v4);
upper.push(2);
g.add_edge(v3, v2);
upper.push(11);
g.add_edge(v3, v4);
upper.push(4);
g.add_edge(v4, v2);
upper.push(4);
g.add_edge(v4, t);
upper.push(20);
let g = g.into_graph();
let (value, flow, _) = pushrelabel(&g, s, t, |e| upper[e.index()]);
assert_eq!(value, 11);
assert!(flow.iter().all(|&(e, f)| f >= 0 && f <= upper[e.index()]));
assert!(g.nodes().filter(|&u| u != s && u != t).all(|u| g
.outedges(u)
.map(|(e, _)| flow[e.index()].1)
.sum::<isize>()
== g.inedges(u).map(|(e, _)| flow[e.index()].1).sum::<isize>()));
}
}
pub fn pushrelabel<'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 = PushRelabel::new(g);
maxflow.solve(src, snk, upper);
(maxflow.value(), maxflow.flow_iter().collect(), maxflow.mincut())
}