competitive/graph/
project_selection_problem.rs

1use super::Dinic;
2use std::{cmp::Ordering, collections::HashMap};
3
4#[derive(Debug, Default, Clone)]
5pub struct ProjectSelectionProblem {
6    n_values: Vec<usize>,
7    start: Vec<usize>,
8    cost1: Vec<Vec<i64>>,
9    cost2: HashMap<(usize, usize), u64>,
10    totalcost: i64,
11}
12impl ProjectSelectionProblem {
13    pub fn new(n_project: usize, n_value: usize) -> Self {
14        Self {
15            n_values: vec![n_value; n_project],
16            start: (0..=n_project * (n_value - 1))
17                .step_by(n_value - 1)
18                .collect(),
19            cost1: vec![vec![0i64; n_value]; n_project],
20            cost2: Default::default(),
21            totalcost: 0i64,
22        }
23    }
24    pub fn with_n_values(n_values: Vec<usize>) -> Self {
25        let mut start = Vec::with_capacity(n_values.len() + 1);
26        start.push(0usize);
27        for nv in n_values.iter() {
28            start.push(start.last().unwrap() + nv - 1);
29        }
30        let cost1 = n_values.iter().map(|&n| vec![0i64; n]).collect();
31        Self {
32            n_values,
33            start,
34            cost1,
35            cost2: Default::default(),
36            totalcost: 0i64,
37        }
38    }
39    pub fn add_cost1(&mut self, p: usize, v: usize, c: i64) {
40        self.cost1[p][v] += c;
41    }
42    /// x1 >= v1 && x2 < v2 (0 < v1 < nv1, 0 < v2 < nv2)
43    pub fn add_cost2_01(&mut self, p1: usize, p2: usize, v1: usize, v2: usize, c: u64) {
44        debug_assert!(0 < v1 && v1 < self.n_values[p1]);
45        debug_assert!(0 < v2 && v2 < self.n_values[p2]);
46        let key = (self.start[p1] + v1 - 1, self.start[p2] + v2 - 1);
47        if c > 0 {
48            *self.cost2.entry(key).or_default() += c;
49        }
50    }
51    /// x1 < v1 && x2 >= v2 (0 < v1 < nv1, 0 < v2 < nv2)
52    pub fn add_cost2_10(&mut self, p1: usize, p2: usize, v1: usize, v2: usize, c: u64) {
53        self.add_cost2_01(p2, p1, v2, v1, c);
54    }
55    /// cost is monge: cost(v1-1, v2) + cost(v1, v2-1) >= cost(v1, v2) + cost(v1-1, v2-1)
56    pub fn add_cost2<F>(&mut self, p1: usize, p2: usize, mut cost: F)
57    where
58        F: FnMut(usize, usize) -> i64,
59    {
60        debug_assert_ne!(p1, p2);
61        let nv1 = self.n_values[p1];
62        let nv2 = self.n_values[p2];
63        debug_assert_ne!(nv1, 0);
64        debug_assert_ne!(nv2, 0);
65        let c00 = cost(0, 0);
66        self.totalcost += c00;
67        for v1 in 1usize..nv1 {
68            self.add_cost1(p1, v1, cost(v1, 0) - c00);
69        }
70        for v2 in 1usize..nv2 {
71            self.add_cost1(p2, v2, cost(0, v2) - c00);
72        }
73        let mut acc = 0i64;
74        for v1 in 1usize..nv1 {
75            for v2 in 1usize..nv2 {
76                let c = cost(v1 - 1, v2) + cost(v1, v2 - 1) - cost(v1, v2) - cost(v1 - 1, v2 - 1);
77                debug_assert!(c >= 0, "cost is not monge");
78                let key = (self.start[p1] + v1 - 1, self.start[p2] + v2 - 1);
79                if c > 0 {
80                    *self.cost2.entry(key).or_default() += c as u64;
81                }
82                acc -= c;
83            }
84            self.add_cost1(p1, v1, acc);
85        }
86    }
87    pub fn solve(&self) -> (i64, Vec<usize>) {
88        let vsize = *self.start.last().unwrap();
89        let esize_expect = vsize * 2 + self.cost2.len();
90        let mut builder = Dinic::builder(vsize + 2, esize_expect);
91        let mut totalcost = self.totalcost;
92        let s = vsize;
93        let t = s + 1;
94        for (p, c) in self.cost1.iter().enumerate() {
95            let nv = self.n_values[p];
96            totalcost += c[nv - 1];
97            for v in 1usize..nv {
98                let u = self.start[p] + v - 1;
99                let d = c[v] - c[v - 1];
100                match d.cmp(&0) {
101                    Ordering::Less => {
102                        builder.add_edge(s, u, (-d) as u64);
103                    }
104                    Ordering::Greater => {
105                        builder.add_edge(u, t, d as u64);
106                        totalcost -= d;
107                    }
108                    Ordering::Equal => {}
109                }
110                if v >= 2 {
111                    builder.add_edge(u, u - 1, u64::MAX);
112                }
113            }
114        }
115        for (&(x, y), &c) in self.cost2.iter() {
116            builder.add_edge(x, y, c);
117        }
118        let dgraph = builder.gen_graph();
119        let mut dinic = builder.build(&dgraph);
120        let res = dinic.maximum_flow(s, t) as i64 + totalcost;
121        let visited = dinic.minimum_cut(s);
122        let mut values = vec![0usize; self.n_values.len()];
123        for (p, &nv) in self.n_values.iter().enumerate() {
124            for v in 1usize..nv {
125                values[p] += visited[self.start[p] + v - 1] as usize;
126            }
127        }
128        (res, values)
129    }
130}
competitive/graph/project_selection_problem.rs

competitive/graph/
project_selection_problem.rs
