competitive/graph/
minimum_cost_flow.rs

1use super::BidirectionalSparseGraph;
2
3#[derive(Debug, Clone)]
4pub struct PrimalDualBuilder {
5    vsize: usize,
6    edges: Vec<(usize, usize)>,
7    capacities: Vec<u64>,
8    costs: Vec<i64>,
9    has_negedge: bool,
10}
11impl PrimalDualBuilder {
12    pub fn new(vsize: usize, esize_expect: usize) -> Self {
13        Self {
14            vsize,
15            edges: Vec::with_capacity(esize_expect),
16            capacities: Vec::with_capacity(esize_expect * 2),
17            costs: Vec::with_capacity(esize_expect * 2),
18            has_negedge: false,
19        }
20    }
21    pub fn add_edge(&mut self, from: usize, to: usize, cap: u64, cost: i64) {
22        self.edges.push((from, to));
23        self.capacities.push(cap);
24        self.capacities.push(0);
25        self.has_negedge |= cost < 0;
26        self.costs.push(cost);
27        self.costs.push(-cost);
28    }
29    pub fn gen_graph(&mut self) -> BidirectionalSparseGraph {
30        let edges = std::mem::take(&mut self.edges);
31        BidirectionalSparseGraph::from_edges(self.vsize, edges)
32    }
33    pub fn build(self, graph: &BidirectionalSparseGraph) -> PrimalDual<'_> {
34        let PrimalDualBuilder {
35            vsize,
36            capacities,
37            costs,
38            has_negedge,
39            ..
40        } = self;
41        PrimalDual {
42            graph,
43            capacities,
44            costs,
45            potential: std::iter::repeat(0).take(vsize).collect(),
46            dist: Vec::with_capacity(vsize),
47            prev_vertex: std::iter::repeat(0).take(vsize).collect(),
48            prev_edge: std::iter::repeat(0).take(vsize).collect(),
49            has_negedge,
50        }
51    }
52}
53impl Extend<(usize, usize, u64, i64)> for PrimalDualBuilder {
54    fn extend<T: IntoIterator<Item = (usize, usize, u64, i64)>>(&mut self, iter: T) {
55        for (from, to, cap, cost) in iter {
56            self.add_edge(from, to, cap, cost)
57        }
58    }
59}
60
61#[derive(Debug)]
62pub struct PrimalDual<'a> {
63    graph: &'a BidirectionalSparseGraph,
64    capacities: Vec<u64>,
65    costs: Vec<i64>,
66    potential: Vec<i64>,
67    dist: Vec<i64>,
68    prev_vertex: Vec<usize>,
69    prev_edge: Vec<usize>,
70    has_negedge: bool,
71}
72impl PrimalDual<'_> {
73    pub fn builder(vsize: usize, esize_expect: usize) -> PrimalDualBuilder {
74        PrimalDualBuilder::new(vsize, esize_expect)
75    }
76    fn bellman_ford(&mut self, s: usize) {
77        self.potential.clear();
78        self.potential.resize(self.graph.vertices_size(), i64::MAX);
79        self.potential[s] = 0;
80        for _ in 1..self.graph.vertices_size() {
81            let mut end = true;
82            for u in self.graph.vertices() {
83                for a in self.graph.adjacencies(u) {
84                    let ncost = self.potential[u].saturating_add(self.costs[a.id]);
85                    if self.capacities[a.id] > 0 && self.potential[a.to] > ncost {
86                        self.potential[a.to] = ncost;
87                        end = false;
88                    }
89                }
90            }
91            if end {
92                break;
93            }
94        }
95    }
96    fn dijkstra(&mut self, s: usize, t: usize) -> bool {
97        use std::{cmp::Reverse, collections::BinaryHeap};
98        self.dist.clear();
99        self.dist.resize(self.graph.vertices_size(), i64::MAX);
100        self.dist[s] = 0;
101        let mut heap = BinaryHeap::new();
102        heap.push((Reverse(0), s));
103        while let Some((Reverse(d), u)) = heap.pop() {
104            if self.dist[u] < d {
105                continue;
106            }
107            if !self.has_negedge && u == t {
108                break;
109            }
110            for a in self.graph.adjacencies(u) {
111                let ncost = (self.dist[u].saturating_add(self.costs[a.id]))
112                    .saturating_add(self.potential[u].saturating_sub(self.potential[a.to]));
113                if self.capacities[a.id] > 0 && self.dist[a.to] > ncost {
114                    self.dist[a.to] = ncost;
115                    self.prev_vertex[a.to] = u;
116                    self.prev_edge[a.to] = a.id;
117                    heap.push((Reverse(ncost), a.to));
118                }
119            }
120        }
121        self.dist[t] != i64::MAX
122    }
123    /// Return (flow, cost).
124    pub fn minimum_cost_flow_limited(&mut self, s: usize, t: usize, limit: u64) -> (u64, i64) {
125        let mut flow = 0;
126        let mut cost = 0;
127        if self.has_negedge {
128            self.bellman_ford(s);
129        }
130        while flow < limit && self.dijkstra(s, t) {
131            for (p, d) in self.potential.iter_mut().zip(self.dist.iter()) {
132                *p = p.saturating_add(*d);
133            }
134            let mut f = limit - flow;
135            let mut v = t;
136            while v != s {
137                f = f.min(self.capacities[self.prev_edge[v]]);
138                v = self.prev_vertex[v];
139            }
140            flow += f;
141            cost += f as i64 * self.potential[t];
142            let mut v = t;
143            while v != s {
144                self.capacities[self.prev_edge[v]] -= f;
145                self.capacities[self.prev_edge[v] ^ 1] += f;
146                v = self.prev_vertex[v];
147            }
148        }
149        (flow, cost)
150    }
151    /// Return (flow, cost).
152    pub fn minimum_cost_flow(&mut self, s: usize, t: usize) -> (u64, i64) {
153        self.minimum_cost_flow_limited(s, t, u64::MAX)
154    }
155    pub fn get_flow(&self, eid: usize) -> u64 {
156        self.capacities[eid * 2 + 1]
157    }
158}