push_relabel 0.1.0

use std::collections::VecDeque;

type Flow = i64;

#[derive(Clone)]
struct Edge {
    to: usize,
    rev: usize,
    f: Flow,
    c: Flow,
}

pub struct PushRelabel {
    n: usize,
    g: Vec<Vec<Edge>>,
    ec: Vec<Flow>,
    cur: Vec<usize>,     // cursor: next edge index to try
    hs: Vec<Vec<usize>>, // active vertices by height
    h: Vec<usize>,       // heights
}

impl PushRelabel {
    pub fn new(n: usize) -> Self {
        Self {
            n,
            g: vec![vec![]; n],
            ec: vec![0; n],
            cur: vec![0; n],
            hs: vec![vec![]; 2 * n],
            h: vec![0; n],
        }
    }

    /// Add an edge u→v with capacity `cap` and a reverse edge t→s with `rcap`.
    /// Self‐loops are ignored.
    pub fn add_edge(&mut self, u: usize, v: usize, cap: Flow, rcap: Flow) {
        if u != v {
            let rev_v = self.g[v].len();
            let rev_u = self.g[u].len();
            self.g[u].push(Edge {
                to: v,
                rev: rev_v,
                f: 0,
                c: cap,
            });
            self.g[v].push(Edge {
                to: u,
                rev: rev_u,
                f: 0,
                c: rcap,
            });
        }
    }

    fn push(&mut self, u: usize, idx: usize, f: Flow) {
        let (v, rev) = (self.g[u][idx].to, self.g[u][idx].rev);
        if self.ec[v] == 0 && f != 0 {
            self.hs[self.h[v]].push(v);
        }
        self.g[u][idx].f += f;
        self.g[u][idx].c -= f;
        self.ec[v] += f;
        self.g[v][rev].f -= f;
        self.g[v][rev].c += f;
        self.ec[u] -= f;
    }

    /// Global relabel: BFS from sink to recompute distance labels and rebuild active lists.
    fn global_relabel(&mut self, s: usize, t: usize, count: &mut [usize]) {
        let n = self.n;
        self.h.fill(n + 1);
        count.fill(0);
        let mut queue = VecDeque::new();
        self.h[t] = 0;
        count[0] = 1;
        queue.push_back(t);
        while let Some(v) = queue.pop_front() {
            let hv = self.h[v];
            for e in &self.g[v] {
                if self.h[e.to] == n + 1 && self.g[e.to][e.rev].c > 0 {
                    self.h[e.to] = hv + 1;
                    count[hv + 1] += 1;
                    queue.push_back(e.to);
                }
            }
        }
        self.h[s] = n;
        count[n] += 1;
        self.cur.fill(0);
        self.hs.iter_mut().for_each(Vec::clear);
        for u in 0..n {
            if u != s && u != t && self.ec[u] > 0 {
                self.hs[self.h[u]].push(u);
            }
        }
    }

    /// Computes max‐flow from `s` to `t` in O(V²√E) with gap heuristic.
    /// Returns the total flow value.
    pub fn max_flow(&mut self, s: usize, t: usize) -> Flow {
        if s == t {
            return 0;
        }
        let n = self.n;
        self.h[s] = n;
        self.ec[t] = 1; // sentinel so it doesn't become active
        let mut count = vec![0; 2 * n];
        count[0] = n - 1;
        self.cur.fill(0);
        for idx in 0..self.g[s].len() {
            self.push(s, idx, self.g[s][idx].c);
        }
        self.global_relabel(s, t, &mut count);
        let mut work = 0;
        let mut hi = self
            .hs
            .iter()
            .rposition(|bucket| !bucket.is_empty())
            .unwrap_or(0);
        loop {
            if self.hs[hi].is_empty() {
                if hi == 0 {
                    break;
                }
                hi -= 1;
                continue;
            }
            let u = self.hs[hi].pop().unwrap();
            while self.ec[u] > 0 {
                if self.cur[u] == self.g[u].len() {
                    count[self.h[u]] -= 1;
                    self.h[u] = usize::MAX;
                    for (i, e) in self.g[u].iter().enumerate() {
                        if e.c > 0 && self.h[u] > self.h[e.to] + 1 {
                            self.h[u] = self.h[e.to] + 1;
                            self.cur[u] = i;
                        }
                    }
                    count[self.h[u]] += 1;
                    hi = self.h[u];
                    if count[self.h[u]] == 0 && self.h[u] < n {
                        for v in 0..n {
                            if self.h[v] > self.h[u] && self.h[v] < n {
                                count[self.h[v]] -= 1;
                                self.h[v] = n + 1;
                            }
                        }
                    }
                    work += 1;
                    if work >= n {
                        self.global_relabel(s, t, &mut count);
                        work = 0;
                        hi = self
                            .hs
                            .iter()
                            .rposition(|bucket| !bucket.is_empty())
                            .unwrap_or(0);
                    }
                } else {
                    let idx = self.cur[u];
                    let e = &self.g[u][idx];
                    if e.c > 0 && self.h[u] == self.h[e.to] + 1 {
                        let send = self.ec[u].min(e.c);
                        self.push(u, idx, send);
                    } else {
                        self.cur[u] += 1;
                    }
                }
            }
        }

        self.ec[t] - 1
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// 1) Trivial: no edges at all → flow = 0
    #[test]
    fn trivial_no_edges() {
        let mut fl = PushRelabel::new(2);
        assert_eq!(fl.max_flow(0, 1), 0);
    }

    /// 2) Single edge
    #[test]
    fn single_edge() {
        let mut fl = PushRelabel::new(2);
        fl.add_edge(0, 1, 7, 0);
        assert_eq!(fl.max_flow(0, 1), 7);
    }

    /// 3) Parallel edges sum properly
    #[test]
    fn parallel_edges() {
        let mut fl = PushRelabel::new(3);
        // two edges 0→1 of caps 3 and 2, then 1→2 cap 5
        fl.add_edge(0, 1, 3, 0);
        fl.add_edge(0, 1, 2, 0);
        fl.add_edge(1, 2, 5, 0);
        assert_eq!(fl.max_flow(0, 2), 5);
    }

    /// 4) Diamond (two disjoint paths)
    #[test]
    fn diamond() {
        let mut fl = PushRelabel::new(5);
        // 0→1→4 and 0→2→4, each cap 10
        fl.add_edge(0, 1, 10, 0);
        fl.add_edge(1, 4, 10, 0);
        fl.add_edge(0, 2, 10, 0);
        fl.add_edge(2, 4, 10, 0);
        // also cross‐link 1→2 and 2→1 small to test back‐edges
        fl.add_edge(1, 2, 1, 0);
        fl.add_edge(2, 1, 1, 0);
        assert_eq!(fl.max_flow(0, 4), 20);
    }

    /// 5) Bottle‐necked multiple paths
    #[test]
    fn bottleneck_paths() {
        let mut fl = PushRelabel::new(6);
        // three parallel paths from 0→5 but with bottlenecks
        fl.add_edge(0, 1, 5, 0);
        fl.add_edge(1, 5, 1, 0);
        fl.add_edge(0, 2, 5, 0);
        fl.add_edge(2, 5, 2, 0);
        fl.add_edge(0, 3, 5, 0);
        fl.add_edge(3, 5, 3, 0);
        // total = 1+2+3 = 6
        assert_eq!(fl.max_flow(0, 5), 6);
    }

    /// 6) Disconnected component + gap heuristic trigger
    #[test]
    fn disconnected_and_gap() {
        let mut fl = PushRelabel::new(4);
        // 0→1 cap=5 but 1→3 cap=0 → no flow via 1
        fl.add_edge(0, 1, 5, 0);
        fl.add_edge(1, 3, 0, 0);
        // 0→2→3 path cap=3
        fl.add_edge(0, 2, 3, 0);
        fl.add_edge(2, 3, 3, 0);

        // we expect only the 0–2–3 path to carry flow
        assert_eq!(fl.max_flow(0, 3), 3);
    }

    /// 7) Unit‐capacity bipartite matching 4×4
    #[test]
    fn bipartite_matching_4x4() {
        let n = 2 + 4 + 4; // source + left + right + sink
        let s = 0;
        let t = n - 1;
        let mut fl = PushRelabel::new(n);

        // source to left
        for u in 0..4 {
            fl.add_edge(s, 1 + u, 1, 0);
        }
        // left to right: connect i→j if (i+j) even
        for i in 0..4 {
            for j in 0..4 {
                if (i + j) % 2 == 0 {
                    fl.add_edge(1 + i, 1 + 4 + j, 1, 0);
                }
            }
        }
        // right to sink
        for v in 0..4 {
            fl.add_edge(1 + 4 + v, t, 1, 0);
        }

        // In this pattern, maximum matching size = 4
        assert_eq!(fl.max_flow(s, t), 4);
    }

    /// 8) Long chain of size 100 with varying capacities
    #[test]
    fn long_chain() {
        let n = 101;
        let mut fl = PushRelabel::new(n);
        // chain 0→1→2→…→100
        for i in 0..100 {
            fl.add_edge(i, i + 1, (i + 1) as Flow, 0);
        }
        // bottleneck is the smallest cap = 1
        assert_eq!(fl.max_flow(0, 100), 1);
    }

    /// 9) Layered grid: 5×5 grid where flow can snake through.
    #[test]
    fn grid_5x5() {
        let w = 5;
        let h = 5;
        // node index = r*w + c, plus s,t at the ends
        let n = w * h + 2;
        let s = w * h;
        let t = w * h + 1;
        let mut fl = PushRelabel::new(n);

        // connect source to top‐row
        for c in 0..w {
            fl.add_edge(s, c, 2, 0);
        }
        // connect bottom‐row to sink
        for c in 0..w {
            fl.add_edge((h - 1) * w + c, t, 2, 0);
        }
        // connect each cell to its 4‐neighbors with cap=1
        let dirs = [(1, 0), (-1, 0), (0, 1), (0, -1)];
        for r in 0..h {
            for c in 0..w {
                for &(dr, dc) in &dirs {
                    let nr = r as isize + dr;
                    let nc = c as isize + dc;
                    if nr >= 0 && nr < h as isize && nc >= 0 && nc < w as isize {
                        let u = r * w + c;
                        let v = (nr as usize) * w + (nc as usize);
                        fl.add_edge(u, v, 1, 0);
                    }
                }
            }
        }

        // Rough lower bound: can send at least w=5 units
        let got = fl.max_flow(s, t);
        assert!(got >= 5 && got <= 10, "unexpected grid flow = {}", got);
    }

    /// 10) Self‐loops and zero‐capacity edges are ignored
    #[test]
    fn self_loops_and_zero_caps() {
        let mut fl = PushRelabel::new(3);
        fl.add_edge(0, 0, 10, 0); // self‐loop
        fl.add_edge(0, 1, 0, 0); // zero cap
        fl.add_edge(0, 2, 5, 0); // valid
        fl.add_edge(2, 1, 5, 0);
        assert_eq!(fl.max_flow(0, 1), 5);
    }

    /// 11) Complete graph small: 4 nodes, all edges cap=1, flow = min‐cut = 3
    #[test]
    fn complete_graph_4() {
        let n = 4;
        let mut fl = PushRelabel::new(n);
        for u in 0..n {
            for v in 0..n {
                if u != v {
                    fl.add_edge(u, v, 1, 0);
                }
            }
        }
        // With node 0 as s and 3 as t, min‐cut(s,t)=3
        assert_eq!(fl.max_flow(0, 3), 3);
    }

    /// 12) Random but deterministic small graph to catch subtle bugs
    #[test]
    fn small_deterministic_random() {
        let mut fl = PushRelabel::new(7);
        // a fixed “random” pattern of edges
        let edges = [
            (0, 1, 3),
            (0, 2, 5),
            (1, 3, 4),
            (2, 3, 2),
            (1, 4, 2),
            (4, 5, 3),
            (5, 3, 1),
            (2, 6, 1),
            (6, 3, 2),
        ];
        for &(u, v, c) in &edges {
            fl.add_edge(u, v, c, 0);
        }
        assert_eq!(fl.max_flow(0, 3), 6);
    }

    /// Helper: build a PushRelabel instance with `n` nodes, add all `edges`,
    /// then compute max_flow(s, t).
    fn max_flow(n: usize, edges: &[(usize, usize, i64, i64)], s: usize, t: usize) -> i64 {
        let mut pr = PushRelabel::new(n);
        for &(u, v, cap, rcap) in edges {
            pr.add_edge(u, v, cap, rcap);
        }
        pr.max_flow(s, t)
    }

    #[test]
    fn empty_single_node() {
        // n=1, source==sink, no edges => flow = 0
        assert_eq!(max_flow(1, &[], 0, 0), 0);
    }

    #[test]
    fn two_nodes_no_edges() {
        // n=2, no path from 0 to 1 => 0
        assert_eq!(max_flow(2, &[], 0, 1), 0);
    }

    #[test]
    fn zero_capacity_edge() {
        // edge of cap=0 should contribute nothing
        let edges = &[(0, 1, 0, 0)];
        assert_eq!(max_flow(2, edges, 0, 1), 0);
    }

    #[test]
    fn simple_chain() {
        // 0→1→2→3 with caps 5,6,7 => bottleneck=5
        let edges = &[(0, 1, 5, 0), (1, 2, 6, 0), (2, 3, 7, 0)];
        assert_eq!(max_flow(4, edges, 0, 3), 5);
    }

    #[test]
    fn diamond_structure() {
        //    0
        //   / \
        //  1   2
        //   \ /
        //    3
        //
        // caps: 0→1=5, 0→2=7, 1→3=4, 2→3=8 => max-flow = 4 + 7 = 11
        let edges = &[(0, 1, 5, 0), (0, 2, 7, 0), (1, 3, 4, 0), (2, 3, 8, 0)];
        assert_eq!(max_flow(4, edges, 0, 3), 11);
    }

    #[test]
    fn disconnected_sink() {
        // 0→1 exists but 1 is not connected to 2 (the sink) => 0
        let edges = &[(0, 1, 5, 0)];
        assert_eq!(max_flow(3, edges, 0, 2), 0);
    }

    #[test]
    fn reverse_capacity_only() {
        // rcap > 0 but forward cap = 0 => no forward flow possible
        let edges = &[(0, 1, 0, 5)];
        assert_eq!(max_flow(2, edges, 0, 1), 0);
    }

    #[test]
    fn small_cycle_with_two_exits() {
        // cycle between 1 and 2, flow can exit to 3:
        // 0→1 cap=10, 1→2 cap=5, 2→1 cap=0, 1→3 cap=3, 2→3 cap=7 => total 3+5=8
        let edges = &[
            (0, 1, 10, 0),
            (1, 2, 5, 0),
            (2, 1, 0, 0),
            (1, 3, 3, 0),
            (2, 3, 7, 0),
        ];
        assert_eq!(max_flow(4, edges, 0, 3), 8);
    }

    #[test]
    fn cross_edge_enhancement() {
        // s=0, a=1, b=2, t=3:
        // 0→1 cap=3, 0→2 cap=2, 1→2 cap=1, 1→3 cap=2, 2→3 cap=3 => expected 5
        let edges = &[
            (0, 1, 3, 0),
            (0, 2, 2, 0),
            (1, 2, 1, 0),
            (1, 3, 2, 0),
            (2, 3, 3, 0),
        ];
        assert_eq!(max_flow(4, edges, 0, 3), 5);
    }

    #[test]
    fn bipartite_matching_k3_3() {
        // K_{3,3} unit‐capacity matching:
        // nodes 1–3 on left, 4–6 on right, super-source=0, super-sink=7
        let mut edges = Vec::new();
        // source to left
        for u in 1..=3 {
            edges.push((0, u, 1, 0));
        }
        // left to right complete
        for u in 1..=3 {
            for v in 4..=6 {
                edges.push((u, v, 1, 0));
            }
        }
        // right to sink
        for v in 4..=6 {
            edges.push((v, 7, 1, 0));
        }
        assert_eq!(max_flow(8, &edges, 0, 7), 3);
    }

    #[test]
    fn layered_network_many_paths() {
        // 6 layers (including source+sink), width=30 → max flow = 30
        let layers = 6;
        let width = 30;
        let mut edges = Vec::new();
        // assign node IDs: layer 0 is source=0.
        let mut next_id = 1;
        let mut layer_nodes = vec![vec![0]];
        for _ in 1..layers {
            let mut this = Vec::with_capacity(width);
            for _ in 0..width {
                this.push(next_id);
                next_id += 1;
            }
            layer_nodes.push(this);
        }
        let sink = next_id;
        let n = sink + 1;
        layer_nodes.push(vec![sink]);
        // fully connect each layer to the next with cap=1
        for l in 0..layers {
            for &u in &layer_nodes[l] {
                for &v in &layer_nodes[l + 1] {
                    edges.push((u, v, 1, 0));
                }
            }
        }
        assert_eq!(max_flow(n, &edges, 0, sink), width as i64);
    }

    #[test]
    fn massive_parallel_edges() {
        // 10 000 parallel unit-cap edges between 0 and 1
        let m = 10_000;
        let edges = vec![(0, 1, 1, 0); m];
        assert_eq!(max_flow(2, &edges, 0, 1), m as i64);
    }

    #[test]
    fn complete_bipartite_k50_50() {
        // standard bipartite matching K₅₀,₅₀
        let left = 50;
        let right = 50;
        let source = 0;
        let sink = left + right + 1;
        let n = sink + 1;
        let mut edges = Vec::new();
        // source → left
        for u in 1..=left {
            edges.push((source, u, 1, 0));
        }
        // left → right (complete)
        for u in 1..=left {
            for v in (left + 1)..=(left + right) {
                edges.push((u, v, 1, 0));
            }
        }
        // right → sink
        for v in (left + 1)..=(left + right) {
            edges.push((v, sink, 1, 0));
        }
        assert_eq!(max_flow(n, &edges, source, sink), left as i64);
    }

    #[test]
    fn dead_end_branches_trigger_gap() {
        // main path 0→1→2→3 (cap=5), dead branch 0→4→5
        let edges = &[
            (0, 1, 5, 0),
            (1, 2, 5, 0),
            (2, 3, 5, 0),
            (0, 4, 10, 0),
            (4, 5, 10, 0),
        ];
        // only the 0→1→2→3 chain reaches sink=3
        assert_eq!(max_flow(6, edges, 0, 3), 5);
    }

    #[test]
    fn decreasing_capacity_chain() {
        // chain of length 10, capacities 100→99→…→91 → bottleneck = 91
        let mut edges = Vec::new();
        for i in 0..10 {
            let cap = 100 - i as i64;
            edges.push((i, i + 1, cap, 0));
        }
        assert_eq!(max_flow(11, &edges, 0, 10), 91);
    }

    /// 25) Reverse capacities do not spuriously increase forward flow.
    #[test]
    fn reverse_capacity_doesnt_add_flow() {
        let mut fl = PushRelabel::new(3);
        // 0→1 cap=5, 1→0 cap=10;  1→2 cap=7, 2→1 cap=3
        fl.add_edge(0, 1, 5, 10);
        fl.add_edge(1, 2, 7, 3);
        // Forward path is limited by 0→1 (5) and 1→2 (7) → min = 5
        assert_eq!(fl.max_flow(0, 2), 5);
    }

    /// 26) Complex cycle should not trap the algorithm or create extra flow.
    #[test]
    fn cycle_flow() {
        let mut fl = PushRelabel::new(4);
        fl.add_edge(0, 1, 8, 0);
        fl.add_edge(1, 2, 3, 0);
        fl.add_edge(2, 0, 5, 0);
        fl.add_edge(2, 3, 4, 0);
        // Only path to 3 is 0→1→2→3 with cap = 3.
        assert_eq!(fl.max_flow(0, 3), 3);
    }

    /// 27) When source == sink, max_flow should always be zero.
    #[test]
    fn source_is_sink() {
        let mut fl = PushRelabel::new(3);
        fl.add_edge(0, 1, 10, 0);
        fl.add_edge(1, 2, 10, 0);
        // Even though edges exist, s == t ⇒ flow = 0
        assert_eq!(fl.max_flow(2, 2), 0);
    }

    /// 28) Star graph: 100 leaves each providing a disjoint unit path.
    #[test]
    fn star_graph() {
        let leaves = 100;
        let n = leaves + 2;
        let s = 0;
        let t = n - 1;
        let mut fl = PushRelabel::new(n);
        for i in 1..=leaves {
            fl.add_edge(s, i, 1, 0);
            fl.add_edge(i, t, 1, 0);
        }
        // Total disjoint unit‐capacity leaves = 100
        assert_eq!(fl.max_flow(s, t), leaves as i64);
    }

    /// 29) Large capacities accumulation to verify no overflow occurs.
    #[test]
    fn large_capacities_sum() {
        let mut fl = PushRelabel::new(2);
        let large = 1_000_000_000_000_i64;
        // Two parallel edges of size 10^12 each
        fl.add_edge(0, 1, large, 0);
        fl.add_edge(0, 1, large, 0);
        assert_eq!(fl.max_flow(0, 1), large * 2);
    }

    /// 30) Many dead ends attached to source to force gap‐relabel on interior nodes.
    #[test]
    fn massive_dead_ends_border() {
        let n = 50;
        let sink = n;
        let mut fl = PushRelabel::new(n + 1);
        // Primary chain 0→1→…→sink with unit capacity
        for i in 0..n {
            fl.add_edge(i, i + 1, 1, 0);
        }
        // Attach 50 “dead‐end” spokes from 0 that never reach the sink
        for i in 0..n {
            fl.add_edge(0, i, 10, 0);
        }
        // Only the main chain can carry flow = 1
        assert_eq!(fl.max_flow(0, sink), 1);
    }

    /// 31) Mixed forward/reverse capacities on a multi‐edge path.
    #[test]
    fn mixed_forward_reverse_capacities() {
        let mut fl = PushRelabel::new(4);
        // 0→1 cap=10, rev=5;  1→2 cap=5, rev=2;  2→3 cap=7, rev=3
        fl.add_edge(0, 1, 10, 5);
        fl.add_edge(1, 2, 5, 2);
        fl.add_edge(2, 3, 7, 3);
        // True max‐flow is limited by the smallest forward cap = 5
        assert_eq!(fl.max_flow(0, 3), 5);
    }
}