haar_lib/flow/
min_cost_flow.rs

1//! 最小費用流
2
3use std::{
4    cmp::{min, Reverse},
5    collections::BinaryHeap,
6};
7
8#[derive(Clone, Debug)]
9struct Edge {
10    to: usize,
11    rev: usize,
12    cap: u64,
13    cost: i64,
14    is_rev: bool,
15}
16
17/// 最小費用流
18#[derive(Clone)]
19pub struct MinCostFlow {
20    size: usize,
21    edges: Vec<Vec<Edge>>,
22}
23
24impl MinCostFlow {
25    /// 頂点数`size`の空の[`MinCostFlow`]を返す。
26    pub fn new(size: usize) -> Self {
27        Self {
28            size,
29            edges: vec![vec![]; size],
30        }
31    }
32
33    /// 頂点`u`から頂点`v`に容量`cap`・費用`cost`の辺を張る。
34    pub fn add_edge(&mut self, u: usize, v: usize, cap: u64, cost: i64) {
35        let rev = self.edges[v].len();
36        self.edges[u].push(Edge {
37            to: v,
38            rev,
39            cap,
40            cost,
41            is_rev: false,
42        });
43        let rev = self.edges[u].len() - 1;
44        self.edges[v].push(Edge {
45            to: u,
46            rev,
47            cap: 0,
48            cost: -cost,
49            is_rev: true,
50        });
51    }
52
53    /// 始点`src`から終点`sink`へ流量`f`を流して、その最小費用を求める。
54    ///
55    /// 流量`f`を流しきれるとき、`Ok`に最小費用を包んで返す。
56    /// そうでないとき、`Err`に流せた流量とその最小費用のペアを包んで返す。
57    pub fn min_cost_flow(&mut self, src: usize, sink: usize, f: u64) -> Result<i64, (u64, i64)> {
58        let mut ret = 0;
59        let mut flow = f;
60        let mut h = vec![0; self.size];
61        let mut prev = vec![(0, 0); self.size];
62        let mut pq = BinaryHeap::<Reverse<(i64, usize)>>::new();
63
64        while flow > 0 {
65            let mut cost = vec![None; self.size];
66
67            cost[src] = Some(0);
68            pq.push(Reverse((0, src)));
69
70            while let Some(Reverse((c, v))) = pq.pop() {
71                if cost[v].unwrap() < c {
72                    continue;
73                }
74
75                for (i, e) in self.edges[v].iter().enumerate() {
76                    if e.cap > 0
77                        && (cost[e.to].is_none()
78                            || cost[e.to].unwrap() + h[e.to] > cost[v].unwrap() + h[v] + e.cost)
79                    {
80                        cost[e.to] = Some(cost[v].unwrap() + e.cost + h[v] - h[e.to]);
81                        prev[e.to] = (v, i);
82                        pq.push(Reverse((cost[e.to].unwrap(), e.to)));
83                    }
84                }
85            }
86
87            if cost[sink].is_none() {
88                break;
89            }
90
91            for i in 0..self.size {
92                if let Some(x) = cost[i] {
93                    h[i] += x;
94                }
95            }
96
97            let mut df = flow;
98            let mut cur = sink;
99            while cur != src {
100                df = min(df, self.edges[prev[cur].0][prev[cur].1].cap);
101                cur = prev[cur].0;
102            }
103
104            flow -= df;
105            ret += df as i64 * h[sink];
106
107            let mut cur = sink;
108            while cur != src {
109                let e = &mut self.edges[prev[cur].0][prev[cur].1];
110                e.cap -= df;
111                let rev = e.rev;
112                self.edges[cur][rev].cap += df;
113                cur = prev[cur].0;
114            }
115        }
116
117        if flow == 0 {
118            Ok(ret)
119        } else {
120            Err((f - flow, ret))
121        }
122    }
123
124    /// 頂点`i`に接続する辺を、`(終点, 流量, 費用)`の形で返す。
125    pub fn get_edges(&self, i: usize) -> Vec<(usize, u64, i64)> {
126        self.edges[i]
127            .iter()
128            .filter(|e| !e.is_rev)
129            .map(|e| (e.to, e.cap, e.cost))
130            .collect()
131    }
132}
haar_lib/flow/min_cost_flow.rs

haar_lib/flow/
min_cost_flow.rs
