Skip to main content

problemreductions/models/graph/
kth_best_spanning_tree.rs

1//! Kth Best Spanning Tree problem implementation.
2//!
3//! Given a weighted graph, determine whether it contains `k` distinct spanning
4//! trees whose total weights are all at most a prescribed bound.
5
6use crate::registry::{CreateSpec, ProblemSchemaEntry, VariantDimension};
7use crate::topology::{Graph, SimpleGraph};
8use crate::traits::Problem;
9use crate::types::WeightElement;
10use num_traits::Zero;
11use serde::{Deserialize, Serialize};
12use std::collections::VecDeque;
13
14inventory::submit! {
15    ProblemSchemaEntry {
16        name: "KthBestSpanningTree",
17        display_name: "Kth Best Spanning Tree",
18        aliases: &[],
19        dimensions: &[VariantDimension::new("weight", "i64", &["i64"])],
20        category: crate::registry::ProblemCategory::Graph,
21        module_path: module_path!(),
22        description: "Do there exist k distinct spanning trees with total weight at most B?",
23        fields: KthBestSpanningTreeCreateSpec::FIELDS,
24    }
25}
26
27/// Kth Best Spanning Tree.
28///
29/// Given an undirected graph `G = (V, E)`, non-negative edge weights `w(e)`,
30/// a positive integer `k`, and a bound `B`, determine whether there are `k`
31/// distinct spanning trees of `G` whose total weights are all at most `B`.
32///
33/// # Representation
34///
35/// A configuration is `k` consecutive binary blocks of length `|E|`.
36/// Each block selects the edges of one candidate spanning tree.
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct KthBestSpanningTree<W: WeightElement> {
39    graph: SimpleGraph,
40    weights: Vec<W>,
41    k: usize,
42    bound: W::Sum,
43}
44
45#[derive(Debug, Deserialize, crate::CreateSpec)]
46struct KthBestSpanningTreeCreateSpec {
47    #[create(codec = "edge-list")]
48    graph: Vec<(usize, usize)>,
49    num_vertices: Option<usize>,
50    #[create(codec = "comma-separated")]
51    edge_weights: Option<Vec<i64>>,
52    k: usize,
53    bound: i64,
54}
55
56impl TryFrom<KthBestSpanningTreeCreateSpec> for KthBestSpanningTree<i64> {
57    type Error = crate::registry::ConstructionError;
58
59    fn try_from(spec: KthBestSpanningTreeCreateSpec) -> Result<Self, Self::Error> {
60        let graph = simple_graph_from_create(spec.graph, spec.num_vertices)?;
61        let weights = spec
62            .edge_weights
63            .unwrap_or_else(|| vec![1; graph.num_edges()]);
64        if weights.len() != graph.num_edges() {
65            return Err(format!(
66                "edge_weights has length {}, expected {}",
67                weights.len(),
68                graph.num_edges()
69            )
70            .into());
71        }
72        if spec.k == 0 {
73            return Err("k must be positive".to_string().into());
74        }
75        Ok(Self::new(graph, weights, spec.k, spec.bound))
76    }
77}
78
79fn simple_graph_from_create(
80    edges: Vec<(usize, usize)>,
81    num_vertices: Option<usize>,
82) -> Result<SimpleGraph, crate::registry::ConstructionError> {
83    if edges.is_empty() && num_vertices.is_none() {
84        return Err("num_vertices is required for an empty graph"
85            .to_string()
86            .into());
87    }
88    for (index, &(u, v)) in edges.iter().enumerate() {
89        if u == v {
90            return Err(format!("graph edge {index} is a self-loop at vertex {u}").into());
91        }
92    }
93    let inferred = edges
94        .iter()
95        .flat_map(|&(u, v)| [u, v])
96        .max()
97        .map(|vertex| vertex.checked_add(1).ok_or("vertex count overflows usize"))
98        .transpose()?
99        .unwrap_or(0);
100    let num_vertices = num_vertices.unwrap_or(inferred);
101    if num_vertices < inferred {
102        return Err(format!("num_vertices {num_vertices} is too small for graph endpoints; need at least {inferred}").into());
103    }
104    Ok(SimpleGraph::new(num_vertices, edges))
105}
106
107impl<W: WeightElement> KthBestSpanningTree<W> {
108    /// Create a new KthBestSpanningTree instance.
109    ///
110    /// # Panics
111    ///
112    /// Panics if the number of weights does not match the number of edges, or
113    /// if `k` is zero.
114    pub fn new(graph: SimpleGraph, weights: Vec<W>, k: usize, bound: W::Sum) -> Self {
115        assert_eq!(
116            weights.len(),
117            graph.num_edges(),
118            "weights length must match graph num_edges"
119        );
120        assert!(k > 0, "k must be positive");
121
122        Self {
123            graph,
124            weights,
125            k,
126            bound,
127        }
128    }
129
130    /// Get the underlying graph.
131    pub fn graph(&self) -> &SimpleGraph {
132        &self.graph
133    }
134
135    /// Get the edge weights.
136    pub fn weights(&self) -> &[W] {
137        &self.weights
138    }
139
140    /// Get the requested number of trees.
141    pub fn k(&self) -> usize {
142        self.k
143    }
144
145    /// Get the weight bound.
146    pub fn bound(&self) -> &W::Sum {
147        &self.bound
148    }
149
150    /// Get the number of vertices.
151    pub fn num_vertices(&self) -> usize {
152        self.graph.num_vertices()
153    }
154
155    /// Get the number of edges.
156    pub fn num_edges(&self) -> usize {
157        self.graph.num_edges()
158    }
159
160    /// Check whether the problem uses a non-unit weight type.
161    pub fn is_weighted(&self) -> bool {
162        !W::IS_UNIT
163    }
164
165    /// Check whether a configuration satisfies the problem.
166    pub fn is_valid_solution(
167        &self,
168        config: &[Vec<bool>],
169    ) -> Result<bool, crate::traits::EvaluationError> {
170        if config.len() != self.k
171            || config
172                .iter()
173                .any(|tree| tree.len() != self.graph.num_edges())
174        {
175            return Ok(false);
176        }
177
178        let edges = self.graph.edges();
179        if !self.blocks_are_pairwise_distinct(config) {
180            return Ok(false);
181        }
182        for tree in config {
183            if !self.block_is_valid_tree(tree, &edges)? {
184                return Ok(false);
185            }
186        }
187        Ok(true)
188    }
189
190    fn block_is_valid_tree(
191        &self,
192        block: &[bool],
193        edges: &[(usize, usize)],
194    ) -> Result<bool, crate::traits::EvaluationError> {
195        if block.len() != edges.len() {
196            return Ok(false);
197        }
198
199        let num_vertices = self.graph.num_vertices();
200        let selected_count = block.iter().filter(|&&selected| selected).count();
201        if selected_count != num_vertices.saturating_sub(1) {
202            return Ok(false);
203        }
204
205        let mut total_weight = W::Sum::zero();
206        let mut adjacency = vec![Vec::new(); num_vertices];
207        let mut start = None;
208
209        for (idx, &selected) in block.iter().enumerate() {
210            if !selected {
211                continue;
212            }
213            total_weight = W::checked_add_to_sum(
214                total_weight,
215                self.weights[idx].to_sum(),
216                "summing spanning tree edge weights",
217            )?;
218            let (u, v) = edges[idx];
219            adjacency[u].push(v);
220            adjacency[v].push(u);
221            if start.is_none() {
222                start = Some(u);
223            }
224        }
225
226        if total_weight > self.bound {
227            return Ok(false);
228        }
229
230        if num_vertices <= 1 {
231            return Ok(true);
232        }
233
234        // SAFETY: num_vertices > 1 and selected_count == num_vertices - 1 > 0,
235        // so at least one edge was selected and `start` is Some.
236        let start = start.expect("at least one selected edge");
237
238        let mut visited = vec![false; num_vertices];
239        let mut queue = VecDeque::new();
240        visited[start] = true;
241        queue.push_back(start);
242
243        while let Some(vertex) = queue.pop_front() {
244            for &neighbor in &adjacency[vertex] {
245                if !visited[neighbor] {
246                    visited[neighbor] = true;
247                    queue.push_back(neighbor);
248                }
249            }
250        }
251
252        Ok(visited.into_iter().all(|seen| seen))
253    }
254
255    fn blocks_are_pairwise_distinct(&self, config: &[Vec<bool>]) -> bool {
256        for left in 0..config.len() {
257            for right in (left + 1)..config.len() {
258                if config[left] == config[right] {
259                    return false;
260                }
261            }
262        }
263        true
264    }
265}
266
267impl<W> Problem for KthBestSpanningTree<W>
268where
269    W: WeightElement + crate::variant::VariantParam,
270{
271    const NAME: &'static str = "KthBestSpanningTree";
272    type Solution = Vec<Vec<bool>>;
273    type Value = crate::types::Or;
274
275    crate::problem_parameters![
276        ("num_vertices", num_vertices),
277        ("num_edges", num_edges),
278        ("k", k),
279    ];
280
281    fn variant() -> Vec<(&'static str, &'static str)> {
282        crate::variant_params![W]
283    }
284
285    fn evaluate(
286        &self,
287        solution: &Self::Solution,
288    ) -> Result<crate::types::Or, crate::traits::EvaluationError> {
289        if solution.len() != self.k
290            || solution
291                .iter()
292                .any(|tree| tree.len() != self.graph.num_edges())
293        {
294            return Err(crate::traits::EvaluationError::InvalidConfiguration(
295                "spanning-tree collection dimensions do not match the instance".into(),
296            ));
297        }
298        Ok(crate::types::Or(self.is_valid_solution(solution)?))
299    }
300}
301
302impl<W> crate::solvers::BruteForceProblem for KthBestSpanningTree<W>
303where
304    W: WeightElement + crate::variant::VariantParam,
305{
306    fn dimensions(&self) -> Vec<usize> {
307        vec![2; self.k * self.graph.num_edges()]
308    }
309}
310
311#[cfg(feature = "example-db")]
312pub(crate) fn canonical_model_example_specs() -> Vec<crate::example_db::specs::ModelExampleSpec> {
313    // K4 with weights [1,1,2,2,2,3], k=2, B=4.
314    // 16 spanning trees; exactly 2 have weight ≤ 4 (both weight 4):
315    //   {01,02,03} (star at 0) and {01,02,13}.
316    // Satisfying configs = 2 (the two orderings of this pair).
317    // 12 variables → 2^12 = 4096 configs, fast to enumerate.
318    let graph = SimpleGraph::new(4, vec![(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]);
319    let problem = KthBestSpanningTree::new(graph, vec![1, 1, 2, 2, 2, 3], 2, 4);
320    vec![crate::example_db::specs::ModelExampleSpec {
321        id: "kth_best_spanning_tree",
322        instance: Box::new(problem),
323        optimal_config: serde_json::json!([
324            [true, true, true, false, false, false],
325            [true, true, false, false, true, false]
326        ]),
327        optimal_value: serde_json::json!(true),
328    }]
329}
330
331crate::declare_variants! {
332    default KthBestSpanningTree<i64> => "2^(num_edges * k)" create KthBestSpanningTreeCreateSpec,
333}
334
335crate::register_brute_force! {
336    KthBestSpanningTree<i64> decode |problem: &KthBestSpanningTree<i64>, indices: Vec<usize>| indices.chunks(problem.num_edges()).map(crate::config::config_to_bits).collect(),
337}
338
339#[cfg(test)]
340#[path = "../../unit_tests/models/graph/kth_best_spanning_tree.rs"]
341mod tests;