1use crate::registry::{CreateSpec, ProblemSchemaEntry, VariantDimension};
8use crate::topology::DirectedGraph;
9use crate::traits::Problem;
10use crate::types::WeightElement;
11use num_traits::Zero;
12use serde::de::Error as _;
13use serde::{Deserialize, Serialize};
14
15inventory::submit! {
16 ProblemSchemaEntry {
17 name: "MultipleChoiceBranching",
18 display_name: "Multiple Choice Branching",
19 aliases: &[],
20 dimensions: &[
21 VariantDimension::new("weight", "i64", &["i64"]),
22 ],
23 category: crate::registry::ProblemCategory::Graph,
24 module_path: module_path!(),
25 description: "Find a branching with partition constraints and weight at least K",
26 fields: MultipleChoiceBranchingCreateSpec::FIELDS,
27 }
28}
29
30#[derive(Debug, Clone, Serialize)]
40pub struct MultipleChoiceBranching<W: WeightElement> {
41 graph: DirectedGraph,
42 weights: Vec<W>,
43 partition: Vec<Vec<usize>>,
44 threshold: W::Sum,
45}
46
47#[derive(Debug, Deserialize, crate::CreateSpec)]
48struct MultipleChoiceBranchingCreateSpec {
49 #[create(codec = "arc-list")]
51 arcs: Vec<(usize, usize)>,
52 num_vertices: Option<usize>,
54 weights: Vec<i64>,
56 partition: Vec<Vec<usize>>,
58 threshold: i64,
60}
61
62impl TryFrom<MultipleChoiceBranchingCreateSpec> for MultipleChoiceBranching<i64> {
63 type Error = crate::registry::ConstructionError;
64 fn try_from(spec: MultipleChoiceBranchingCreateSpec) -> Result<Self, Self::Error> {
65 let inferred = spec
66 .arcs
67 .iter()
68 .flat_map(|&(u, v)| [u, v])
69 .max()
70 .map(|vertex| vertex.checked_add(1).ok_or("vertex count overflows usize"))
71 .transpose()?
72 .unwrap_or(0);
73 let num_vertices = spec.num_vertices.unwrap_or(inferred);
74 if num_vertices < inferred {
75 return Err("num_vertices is too small for arc endpoints"
76 .to_string()
77 .into());
78 }
79 let graph = DirectedGraph::new(num_vertices, spec.arcs);
80 let num_arcs = graph.num_arcs();
81 if spec.weights.len() != num_arcs {
82 return Err(format!(
83 "weights has {} entries, expected {num_arcs}",
84 spec.weights.len()
85 )
86 .into());
87 }
88 if let Some(message) = partition_validation_error(&spec.partition, num_arcs) {
89 return Err(message.into());
90 }
91 Ok(Self::new(
92 graph,
93 spec.weights,
94 spec.partition,
95 spec.threshold,
96 ))
97 }
98}
99
100#[derive(Debug, Deserialize)]
101struct MultipleChoiceBranchingUnchecked<W: WeightElement> {
102 graph: DirectedGraph,
103 weights: Vec<W>,
104 partition: Vec<Vec<usize>>,
105 threshold: W::Sum,
106}
107
108impl<'de, W> Deserialize<'de> for MultipleChoiceBranching<W>
109where
110 W: WeightElement + Deserialize<'de>,
111 W::Sum: Deserialize<'de>,
112{
113 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
114 where
115 D: serde::Deserializer<'de>,
116 {
117 let unchecked = MultipleChoiceBranchingUnchecked::<W>::deserialize(deserializer)?;
118 let num_arcs = unchecked.graph.num_arcs();
119 if unchecked.weights.len() != num_arcs {
120 return Err(D::Error::custom(format!(
121 "weights length must match graph num_arcs (expected {num_arcs}, got {})",
122 unchecked.weights.len()
123 )));
124 }
125 if let Some(message) = partition_validation_error(&unchecked.partition, num_arcs) {
126 return Err(D::Error::custom(message));
127 }
128
129 Ok(Self {
130 graph: unchecked.graph,
131 weights: unchecked.weights,
132 partition: unchecked.partition,
133 threshold: unchecked.threshold,
134 })
135 }
136}
137
138impl<W: WeightElement> MultipleChoiceBranching<W> {
139 pub fn new(
141 graph: DirectedGraph,
142 weights: Vec<W>,
143 partition: Vec<Vec<usize>>,
144 threshold: W::Sum,
145 ) -> Self {
146 let num_arcs = graph.num_arcs();
147 assert_eq!(
148 weights.len(),
149 num_arcs,
150 "weights length must match graph num_arcs"
151 );
152 validate_partition(&partition, num_arcs);
153 Self {
154 graph,
155 weights,
156 partition,
157 threshold,
158 }
159 }
160
161 pub fn graph(&self) -> &DirectedGraph {
163 &self.graph
164 }
165
166 pub fn weights(&self) -> &[W] {
168 &self.weights
169 }
170
171 pub fn set_weights(&mut self, weights: Vec<W>) {
173 assert_eq!(
174 weights.len(),
175 self.graph.num_arcs(),
176 "weights length must match graph num_arcs"
177 );
178 self.weights = weights;
179 }
180
181 pub fn is_weighted(&self) -> bool {
183 !W::IS_UNIT
184 }
185
186 pub fn partition(&self) -> &[Vec<usize>] {
188 &self.partition
189 }
190
191 pub fn threshold(&self) -> &W::Sum {
193 &self.threshold
194 }
195
196 pub fn max_weight_bits(&self) -> u64 {
198 crate::types::max_numeric_magnitude_bits(
199 self.weights
200 .iter()
201 .map(|weight| weight.to_sum())
202 .chain(std::iter::once(self.threshold.clone())),
203 )
204 }
205
206 pub fn num_vertices(&self) -> usize {
208 self.graph.num_vertices()
209 }
210
211 pub fn num_arcs(&self) -> usize {
213 self.graph.num_arcs()
214 }
215
216 pub fn num_partition_groups(&self) -> usize {
218 self.partition.len()
219 }
220
221 pub fn is_valid_solution(
223 &self,
224 config: &[bool],
225 ) -> Result<bool, crate::traits::EvaluationError> {
226 is_valid_multiple_choice_branching(
227 &self.graph,
228 &self.weights,
229 &self.partition,
230 &self.threshold,
231 config,
232 )
233 }
234}
235
236impl<W> Problem for MultipleChoiceBranching<W>
237where
238 W: WeightElement + crate::variant::VariantParam,
239{
240 const NAME: &'static str = "MultipleChoiceBranching";
241 type Solution = Vec<bool>;
242 type Value = crate::types::Or;
243
244 crate::problem_parameters![
245 ("max_weight_bits", max_weight_bits),
246 ("num_vertices", num_vertices),
247 ("num_arcs", num_arcs),
248 ("num_partition_groups", num_partition_groups),
249 ];
250
251 fn variant() -> Vec<(&'static str, &'static str)> {
252 crate::variant_params![W]
253 }
254
255 fn evaluate(
256 &self,
257 config: &Self::Solution,
258 ) -> Result<crate::types::Or, crate::traits::EvaluationError> {
259 if config.len() != self.graph.num_arcs() {
260 return Err(crate::traits::EvaluationError::InvalidConfiguration(
261 "arc-selection length does not match the graph".into(),
262 ));
263 }
264 Ok({
265 crate::types::Or({
266 is_valid_multiple_choice_branching(
267 &self.graph,
268 &self.weights,
269 &self.partition,
270 &self.threshold,
271 config,
272 )?
273 })
274 })
275 }
276}
277
278impl<W> crate::solvers::BruteForceProblem for MultipleChoiceBranching<W>
279where
280 W: WeightElement + crate::variant::VariantParam,
281{
282 fn dimensions(&self) -> Vec<usize> {
283 vec![2; self.graph.num_arcs()]
284 }
285}
286
287fn validate_partition(partition: &[Vec<usize>], num_arcs: usize) {
288 if let Some(message) = partition_validation_error(partition, num_arcs) {
289 panic!("{message}");
290 }
291}
292
293fn partition_validation_error(partition: &[Vec<usize>], num_arcs: usize) -> Option<String> {
294 let mut seen = vec![false; num_arcs];
295 for group in partition {
296 for &arc_index in group {
297 if arc_index >= num_arcs {
298 return Some(format!(
299 "partition arc index {} out of range for {} arcs",
300 arc_index, num_arcs
301 ));
302 }
303 if seen[arc_index] {
304 return Some(format!(
305 "partition arc index {} appears more than once",
306 arc_index
307 ));
308 }
309 seen[arc_index] = true;
310 }
311 }
312 if seen.iter().all(|present| *present) {
313 None
314 } else {
315 Some("partition must cover every arc exactly once".to_string())
316 }
317}
318
319fn is_valid_multiple_choice_branching<W: WeightElement>(
320 graph: &DirectedGraph,
321 weights: &[W],
322 partition: &[Vec<usize>],
323 threshold: &W::Sum,
324 config: &[bool],
325) -> Result<bool, crate::traits::EvaluationError> {
326 if config.len() != graph.num_arcs() {
327 return Ok(false);
328 }
329 for group in partition {
330 if group.iter().filter(|&&arc_index| config[arc_index]).count() > 1 {
331 return Ok(false);
332 }
333 }
334
335 let arcs = graph.arcs();
336 let mut in_degree = vec![0usize; graph.num_vertices()];
337 let mut selected_successors = vec![Vec::new(); graph.num_vertices()];
338 let mut total = W::Sum::zero();
339 for (index, &selected) in config.iter().enumerate() {
340 if selected {
341 let (source, target) = arcs[index];
342 in_degree[target] += 1;
343 if in_degree[target] > 1 {
344 return Ok(false);
345 }
346 selected_successors[source].push(target);
347 total = W::checked_add_to_sum(
348 total,
349 weights[index].to_sum(),
350 "summing multiple-choice branching weights",
351 )?;
352 }
353 }
354
355 if total < *threshold {
356 return Ok(false);
357 }
358
359 let mut queue: Vec<usize> = (0..graph.num_vertices())
360 .filter(|&vertex| in_degree[vertex] == 0)
361 .collect();
362 let mut visited = 0usize;
363 while let Some(source) = queue.pop() {
364 visited += 1;
365 for &target in &selected_successors[source] {
366 in_degree[target] -= 1;
367 if in_degree[target] == 0 {
368 queue.push(target);
369 }
370 }
371 }
372
373 Ok(visited == graph.num_vertices())
374}
375
376crate::declare_variants! {
377 default MultipleChoiceBranching<i64> => "2^num_arcs" create MultipleChoiceBranchingCreateSpec,
378}
379
380crate::register_brute_force! {
381 MultipleChoiceBranching<i64> decode |_, indices: Vec<usize>| crate::config::config_to_bits(&indices),
382}
383
384#[cfg(feature = "example-db")]
385pub(crate) fn canonical_model_example_specs() -> Vec<crate::example_db::specs::ModelExampleSpec> {
386 vec![crate::example_db::specs::ModelExampleSpec {
387 id: "multiple_choice_branching",
388 instance: Box::new(MultipleChoiceBranching::new(
389 DirectedGraph::new(
390 6,
391 vec![
392 (0, 1),
393 (0, 2),
394 (1, 3),
395 (2, 3),
396 (1, 4),
397 (3, 5),
398 (4, 5),
399 (2, 4),
400 ],
401 ),
402 vec![3, 2, 4, 1, 2, 3, 1, 3],
403 vec![vec![0, 1], vec![2, 3], vec![4, 7], vec![5, 6]],
404 10,
405 )),
406 optimal_config: serde_json::json!(vec![true, false, true, false, false, true, false, true]),
407 optimal_value: serde_json::json!(true),
408 }]
409}
410
411#[cfg(test)]
412#[path = "../../unit_tests/models/graph/multiple_choice_branching.rs"]
413mod tests;