Skip to main content

arboretum_td/
tree_decomposition.rs

1use crate::datastructures::BitSet;
2use crate::graph::BaseGraph;
3use fxhash::{FxHashMap, FxHashSet};
4use std::cmp::max;
5use std::fmt;
6use std::fmt::{Display, Formatter};
7
8pub enum TreeDecompositionValidationError {
9    HasCycle,
10    NotConnected,
11    MissingVertex(usize),
12    MissingEdge((usize, usize)),
13    NotInducingSubtree(usize),
14}
15
16impl Display for TreeDecompositionValidationError {
17    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
18        match *self {
19            TreeDecompositionValidationError::HasCycle => write!(f, "Has Cycle"),
20            TreeDecompositionValidationError::NotConnected => write!(f, "Not Connected"),
21            TreeDecompositionValidationError::MissingVertex(v) => {
22                write!(f, "Missing Vertex: {}", v)
23            }
24            TreeDecompositionValidationError::MissingEdge((u, v)) => {
25                write!(f, "Missing Edge: ({}, {})", u, v)
26            }
27            TreeDecompositionValidationError::NotInducingSubtree(v) => {
28                write!(f, "Not Inducing Subtree: {}", v)
29            }
30        }
31    }
32}
33
34#[derive(Debug, Clone, Default)]
35pub struct TreeDecomposition {
36    pub bags: Vec<Bag>,
37    pub root: Option<usize>,
38    pub max_bag_size: usize,
39}
40
41impl TreeDecomposition {
42    pub fn flatten(&mut self) {
43        while let Some((parent, child)) = self.find_combinable() {
44            self.reroute(child, parent);
45            self.remove_bag(child);
46        }
47    }
48
49    pub fn with_root(vertex_set: FxHashSet<usize>) -> Self {
50        let mut td = Self::default();
51        td.add_bag(vertex_set);
52        td
53    }
54
55    fn reroute(&mut self, old_bag: usize, parent_idx: usize) {
56        let old_neighbors = self.bags[old_bag].neighbors.clone();
57
58        self.bags[parent_idx].neighbors.extend(old_neighbors.iter());
59        self.bags[parent_idx].neighbors.remove(&parent_idx);
60        self.bags[parent_idx].neighbors.remove(&old_bag);
61
62        let old_id = self.bags[old_bag].id;
63        for neighbor_idx in old_neighbors {
64            if neighbor_idx == parent_idx {
65                continue;
66            }
67            assert!(!self.bags[neighbor_idx].neighbors.contains(&parent_idx));
68
69            assert!(self.bags[neighbor_idx].neighbors.remove(&old_id));
70            assert!(self.bags[neighbor_idx].neighbors.insert(parent_idx));
71        }
72
73        self.bags[old_bag].neighbors.clear();
74    }
75
76    fn remove_bag(&mut self, id: usize) {
77        assert!(self.bags[id].neighbors.is_empty());
78        if id == self.bags.len() - 1 {
79            //println!(" Only need to pop!");
80            self.bags.pop();
81        } else {
82            //println!(" Pre swap {} {:?}", id, self.bags[id]);
83            let old_last = self.bags.swap_remove(id);
84            //println!(" Post swap {} {:?}", id, self.bags[id]);
85            assert!(old_last.neighbors.is_empty());
86            self.bags[id].id = id;
87            let old_last = self.bags.len();
88            for neighbor in self.bags[id].neighbors.clone() {
89                //println!(" Neighbor {} {:?}", neighbor, self.bags[neighbor]);
90                assert!(self.bags[neighbor].neighbors.remove(&old_last));
91                assert!(self.bags[neighbor].neighbors.insert(id));
92            }
93        }
94    }
95
96    fn find_combinable(&self) -> Option<(usize, usize)> {
97        for b in &self.bags {
98            if let Some(n) = b
99                .neighbors
100                .iter()
101                .find(|n| self.bags[**n].vertex_set.is_subset(&b.vertex_set))
102            {
103                return Some((b.id, self.bags[*n].id));
104            }
105        }
106        None
107    }
108
109    pub fn add_bag(&mut self, vertex_set: FxHashSet<usize>) -> usize {
110        let id = self.bags.len();
111        if id == 0 {
112            self.root = Some(id);
113        }
114        self.max_bag_size = max(self.max_bag_size, vertex_set.len());
115        self.bags.push(Bag {
116            id,
117            vertex_set,
118            neighbors: FxHashSet::default(),
119        });
120        id
121    }
122
123    pub fn add_child_bags(&mut self, parent: usize, children: Vec<FxHashSet<usize>>) -> Vec<usize> {
124        assert!(self.bags.len() < parent);
125        let mut ids = Vec::with_capacity(children.len());
126        for c in children {
127            let id = self.add_bag(c);
128            ids.push(id);
129            self.add_edge(parent, id);
130        }
131        ids
132    }
133
134    pub fn add_edge(&mut self, b1: usize, b2: usize) {
135        assert!(b1 < self.bags.len());
136        assert!(b2 < self.bags.len());
137        assert_ne!(b1, b2);
138        self.bags[b1].neighbors.insert(b2);
139        self.bags[b2].neighbors.insert(b1);
140    }
141
142    pub fn dfs(&self) -> TreeDecompositionIterator {
143        let mut visited = BitSet::new(self.bags.len());
144        let stack = if self.root.is_some() {
145            visited.set_bit(self.root.unwrap());
146            vec![self.root.unwrap()]
147        } else {
148            vec![]
149        };
150        TreeDecompositionIterator {
151            td: self,
152            stack,
153            visited,
154        }
155    }
156
157    pub fn replace_bag_v2(&mut self, target_bag: usize, mut td: TreeDecomposition) {
158        let mut separators: FxHashMap<usize, FxHashSet<usize>> = FxHashMap::default();
159        for neighbor in &self.bags[target_bag].neighbors {
160            let key = self.bags[*neighbor].id;
161            let value: FxHashSet<_> = self.bags[target_bag]
162                .vertex_set
163                .intersection(&self.bags[*neighbor].vertex_set)
164                .copied()
165                .collect();
166            separators.insert(key, value);
167        }
168        let offset = self.bags.len();
169
170        for bag in &mut td.bags {
171            bag.id += offset;
172            bag.neighbors = bag.neighbors.iter().map(|i| *i + offset).collect();
173        }
174
175        self.bags.reserve(td.bags.len());
176        for bag in td.bags {
177            self.bags.push(bag)
178        }
179
180        for (_, separator) in separators {
181            let new_neighbor = self.bags[offset..]
182                .iter_mut()
183                .find(|b| b.vertex_set.is_superset(&separator))
184                .unwrap();
185
186            assert!(new_neighbor.neighbors.remove(&target_bag));
187            assert!(new_neighbor.neighbors.insert(new_neighbor.id));
188        }
189        self.bags.swap_remove(target_bag);
190        self.max_bag_size = self.bags.iter().map(|b| b.vertex_set.len()).max().unwrap();
191    }
192
193    pub fn replace_bag(&mut self, target_bag: usize, mut td: TreeDecomposition) {
194        let neighbors_of_target_bag = self.bags[target_bag].neighbors.clone();
195        let vertices_of_target_bag = self.bags[target_bag].vertex_set.clone();
196        let offset = self.bags.len();
197        td.bags.iter_mut().for_each(|b| {
198            b.id += offset;
199            b.neighbors = b.neighbors.iter().map(|i| *i + offset).collect();
200        });
201        for neighbor_of_target_bag in neighbors_of_target_bag {
202            let neighbor_of_target_bag = &mut self.bags[neighbor_of_target_bag];
203            let intersection: FxHashSet<_> = vertices_of_target_bag
204                .intersection(&neighbor_of_target_bag.vertex_set)
205                .copied()
206                .collect();
207            let new_neighbor_of_neighbor_of_target_bag = td
208                .bags
209                .iter_mut()
210                .find(|b| b.vertex_set.is_superset(&intersection))
211                .unwrap();
212            new_neighbor_of_neighbor_of_target_bag
213                .neighbors
214                .insert(neighbor_of_target_bag.id);
215
216            assert!(neighbor_of_target_bag.neighbors.remove(&target_bag));
217            assert!(neighbor_of_target_bag
218                .neighbors
219                .insert(new_neighbor_of_neighbor_of_target_bag.id));
220        }
221        self.bags.extend_from_slice(&td.bags);
222        let old_idx = self.bags.len() - 1;
223        self.bags.swap(target_bag, old_idx);
224        self.bags[target_bag].id = target_bag;
225        for id in self.bags[target_bag].neighbors.clone() {
226            assert!(self.bags[id].neighbors.remove(&old_idx));
227            assert!(self.bags[id].neighbors.insert(target_bag));
228        }
229        self.bags.swap_remove(old_idx);
230        self.max_bag_size = self.bags.iter().map(|b| b.vertex_set.len()).max().unwrap();
231    }
232
233    pub fn bags(&self) -> &[Bag] {
234        &self.bags
235    }
236
237    pub fn combine_with_or_replace(&mut self, glue_point: usize, other: TreeDecomposition) {
238        if self.bags.is_empty() {
239            *self = other;
240        } else {
241            self.combine_with(glue_point, other);
242        }
243    }
244
245    pub fn combine_with(&mut self, glue_point: usize, mut other: TreeDecomposition) {
246        assert!(glue_point < self.bags.len());
247        self.max_bag_size = max(self.max_bag_size, other.max_bag_size);
248        let offset = self.bags.len();
249        for mut b in other.bags.iter_mut() {
250            b.id += offset;
251            b.neighbors = b.neighbors.iter().map(|n| *n + offset).collect();
252        }
253        let other_glue_point = other
254            .bags
255            .iter_mut()
256            .find(|b| b.vertex_set.is_superset(&self.bags[glue_point].vertex_set))
257            .unwrap();
258        other_glue_point.neighbors.insert(glue_point);
259        self.bags[glue_point].neighbors.insert(other_glue_point.id);
260        self.bags.append(&mut other.bags);
261    }
262
263    pub fn verify<G: BaseGraph>(&self, graph: &G) -> Result<(), TreeDecompositionValidationError> {
264        if !self.is_connected() {
265            return Err(TreeDecompositionValidationError::NotConnected);
266        }
267
268        if self.is_cyclic() {
269            return Err(TreeDecompositionValidationError::HasCycle);
270        }
271
272        if let Some(v) = self.get_missing_vertex(graph) {
273            return Err(TreeDecompositionValidationError::MissingVertex(v));
274        }
275
276        if let Some(e) = self.get_missing_edge(graph) {
277            return Err(TreeDecompositionValidationError::MissingEdge(e));
278        }
279
280        if let Some(v) = self.get_vertex_not_inducing_subtree(graph) {
281            return Err(TreeDecompositionValidationError::NotInducingSubtree(v));
282        }
283
284        Ok(())
285    }
286
287    fn is_connected(&self) -> bool {
288        if self.bags.is_empty() {
289            return true;
290        }
291        let mut visited = BitSet::new(self.bags.len());
292        self.dfs().for_each(|b| {
293            visited.set_bit(b.id);
294        });
295        visited.full()
296    }
297
298    fn is_cyclic(&self) -> bool {
299        if self.bags.is_empty() {
300            return true;
301        }
302        let mut visited = BitSet::new(self.bags.len());
303        self.is_cyclic_rec(&mut visited, self.root.unwrap(), None)
304    }
305
306    fn is_cyclic_rec(&self, visited: &mut BitSet, v: usize, parent: Option<usize>) -> bool {
307        visited.set_bit(v);
308        for n in self.bags[v].neighbors.iter().copied() {
309            if !visited[n] {
310                if self.is_cyclic_rec(visited, n, Some(v)) {
311                    return true;
312                }
313            } else if parent.is_some() && n != parent.unwrap() {
314                return true;
315            }
316        }
317        false
318    }
319
320    fn get_missing_vertex<G: BaseGraph>(&self, graph: &G) -> Option<usize> {
321        let mut vertices: FxHashSet<usize> = graph.vertices().collect();
322        self.bags.iter().for_each(|b| {
323            b.vertex_set.iter().for_each(|x| {
324                vertices.remove(x);
325            })
326        });
327        if vertices.is_empty() {
328            None
329        } else {
330            Some(*vertices.iter().next().unwrap())
331        }
332    }
333
334    fn get_missing_edge<G: BaseGraph>(&self, graph: &G) -> Option<(usize, usize)> {
335        for u in graph.vertices() {
336            for v in graph.vertices() {
337                if u < v
338                    && graph.has_edge(u, v)
339                    && !self
340                        .bags
341                        .iter()
342                        .any(|b| b.vertex_set.contains(&u) && b.vertex_set.contains(&v))
343                {
344                    return Some((u, v));
345                }
346            }
347        }
348        None
349    }
350
351    fn get_vertex_not_inducing_subtree<G: BaseGraph>(&self, graph: &G) -> Option<usize> {
352        for u in graph.vertices() {
353            let mut inducing_bags: FxHashSet<usize> = self
354                .bags
355                .iter()
356                .filter(|b| b.vertex_set.contains(&u))
357                .map(|b| b.id)
358                .collect();
359
360            let first = *inducing_bags.iter().next().unwrap();
361            inducing_bags.remove(&first);
362            let mut visited = BitSet::new(self.bags.len());
363            visited.set_bit(first);
364            let mut stack: Vec<usize> = vec![first];
365            while let Some(c) = stack.pop() {
366                for n in self.bags[c].neighbors.iter().copied() {
367                    let bag = &self.bags[n];
368                    if !visited[n] && bag.vertex_set.contains(&u) {
369                        inducing_bags.remove(&bag.id);
370                        stack.push(n);
371                        visited.set_bit(n);
372                    }
373                }
374            }
375            if !inducing_bags.is_empty() {
376                return Some(u + 1);
377            }
378        }
379        None
380    }
381}
382
383pub struct TreeDecompositionIterator<'a> {
384    td: &'a TreeDecomposition,
385    stack: Vec<usize>,
386    visited: BitSet,
387}
388
389impl<'a> Iterator for TreeDecompositionIterator<'a> {
390    type Item = &'a Bag;
391
392    fn next(&mut self) -> Option<Self::Item> {
393        let current = self.stack.pop()?;
394        for c in self.td.bags[current].neighbors.iter().copied() {
395            if !self.visited[c] {
396                self.stack.push(c);
397                self.visited.set_bit(c);
398            }
399        }
400        Some(self.td.bags.get(current).unwrap())
401    }
402}
403
404#[derive(Debug, Default, Clone)]
405pub struct Bag {
406    pub id: usize,
407    pub vertex_set: FxHashSet<usize>,
408    pub neighbors: FxHashSet<usize>,
409}