logicng/datastructures/
ubtree.rs

1use std::collections::{BTreeMap, BTreeSet, HashSet};
2
3type NodeIndex = usize;
4type ChildrenIndex = usize;
5
6struct UbNode<T: Ord> {
7    children: ChildrenIndex,
8    set: Option<BTreeSet<T>>,
9}
10
11impl<T: Ord> UbNode<T> {
12    const fn new(children: ChildrenIndex) -> Self {
13        Self { children, set: None }
14    }
15
16    const fn end_of_path(&self) -> Option<&BTreeSet<T>> {
17        self.set.as_ref()
18    }
19
20    #[cfg(test)]
21    const fn is_end_of_path(&self) -> bool {
22        self.set.is_some()
23    }
24}
25
26/// A data structure for storing sets with efficient sub- and superset queries.
27///
28/// C.f. `A New Method to Index and Query Sets`, Hoffmann and Koehler, 1999.
29/// Note that, we've only implemented the parts of an `UbTree` that we actually
30/// use.
31pub struct UbTree<T: Ord> {
32    root_nodes: ChildrenIndex,
33    nodes: Vec<UbNode<T>>,
34    children_maps: Vec<BTreeMap<T, NodeIndex>>,
35}
36
37impl<T: Ord + Copy> UbTree<T> {
38    /// Construct an empty `UbTree`.
39    pub fn new() -> Self {
40        Self { root_nodes: 0, nodes: Vec::new(), children_maps: vec![BTreeMap::new()] }
41    }
42
43    /// Adds a set of comparable objects to this `UbTree`.
44    pub fn add_set(&mut self, set: BTreeSet<T>) {
45        let mut nodes = self.root_nodes;
46        let mut node = None;
47        for element in &set {
48            let n = if self.children_maps[nodes].contains_key(element) {
49                *self.children_maps[nodes].get(element).unwrap()
50            } else {
51                let new_node = self.new_node();
52                self.children_maps[nodes].insert(*element, new_node);
53                new_node
54            };
55            nodes = self.nodes[n].children;
56            node = Some(n);
57        }
58        if let Some(n) = node {
59            self.nodes[n].set = Some(set);
60        }
61    }
62
63    /// Returns the first subset of a given set in this `UbTree`.
64    pub fn first_subset(&self, set: &BTreeSet<T>) -> Option<&BTreeSet<T>> {
65        if self.get_children(self.root_nodes).is_empty() || set.is_empty() {
66            None
67        } else {
68            self.first_subset_from_forest(set, self.root_nodes)
69        }
70    }
71
72    /// Returns all sets in this `UbTree`.
73    pub fn all_sets(&self) -> BTreeSet<&BTreeSet<T>> {
74        let mut all_sets = BTreeSet::new();
75        let mut stack = Vec::new();
76        stack.push(self.root_nodes);
77        while let Some(current) = stack.pop() {
78            for node_index in self.get_children(current).values() {
79                let node = self.get_node(*node_index);
80                if let Some(set) = node.end_of_path() {
81                    all_sets.insert(set);
82                }
83                stack.push(node.children);
84            }
85        }
86        all_sets
87    }
88
89    #[cfg(test)]
90    pub fn root_nodes(&self) -> &BTreeMap<T, NodeIndex> {
91        self.get_children(self.root_nodes)
92    }
93
94    fn first_subset_from_forest(&self, set: &BTreeSet<T>, forest: ChildrenIndex) -> Option<&BTreeSet<T>> {
95        let nodes = self.get_all_nodes_containing_elements(set, forest);
96        let mut found_subset = None;
97        for node in nodes {
98            if found_subset.is_some() {
99                return found_subset;
100            }
101            if let Some(set) = self.get_node(node).end_of_path() {
102                return Some(set);
103            }
104            let mut remaining_set = set.clone();
105            if let Some(f) = set.first() {
106                remaining_set.remove(f);
107            }
108            found_subset = self.first_subset_from_forest(&remaining_set, self.nodes[node].children);
109        }
110        found_subset
111    }
112
113    fn get_all_nodes_containing_elements(&self, set: &BTreeSet<T>, forest: ChildrenIndex) -> HashSet<NodeIndex> {
114        let mut nodes = HashSet::new();
115        for element in set {
116            let node_index = self.get_children(forest).get(element);
117            if let Some(n) = node_index {
118                nodes.insert(*n);
119            }
120        }
121        nodes
122    }
123
124    fn new_node(&mut self) -> NodeIndex {
125        let node_index = self.nodes.len();
126        let children_index = self.new_children_map(BTreeMap::new());
127        self.nodes.push(UbNode::new(children_index));
128        node_index
129    }
130
131    fn new_children_map(&mut self, map: BTreeMap<T, NodeIndex>) -> ChildrenIndex {
132        let index = self.children_maps.len();
133        self.children_maps.push(map);
134        index
135    }
136
137    fn get_children(&self, index: ChildrenIndex) -> &BTreeMap<T, NodeIndex> {
138        &self.children_maps[index]
139    }
140
141    fn get_node(&self, index: NodeIndex) -> &UbNode<T> {
142        &self.nodes[index]
143    }
144}
145
146impl<T: Ord + Copy> Default for UbTree<T> {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use std::collections::BTreeSet;
155
156    use itertools::Itertools;
157
158    use super::UbTree;
159
160    #[test]
161    fn test_empty_set() {
162        let mut tree = UbTree::<usize>::new();
163        tree.add_set(BTreeSet::new());
164        assert!(tree.root_nodes().is_empty());
165    }
166
167    #[test]
168    fn test_single_set() {
169        let mut tree = UbTree::<usize>::new();
170        tree.add_set(BTreeSet::from_iter([111, 222, 333]));
171        assert_eq!(tree.root_nodes().len(), 1);
172        assert_eq!(tree.root_nodes().keys().copied().collect_vec(), vec![111]);
173
174        let root_a = tree.get_node(tree.root_nodes()[&111]);
175        assert!(!root_a.is_end_of_path());
176        assert_eq!(tree.get_children(root_a.children).len(), 1);
177        assert!(tree.get_children(root_a.children).contains_key(&222));
178
179        let root_a_b = tree.get_node(tree.get_children(root_a.children)[&222]);
180        assert!(!root_a_b.is_end_of_path());
181        assert_eq!(tree.get_children(root_a_b.children).len(), 1);
182        assert!(tree.get_children(root_a_b.children).contains_key(&333));
183
184        let root_a_b_c = tree.get_node(tree.get_children(root_a_b.children)[&333]);
185        assert!(root_a_b_c.is_end_of_path());
186        assert!(tree.get_children(root_a_b_c.children).is_empty());
187        assert_eq!(root_a_b_c.set, Some(BTreeSet::from_iter([111, 222, 333])));
188    }
189
190    #[test]
191    fn test_example_from_paper() {
192        let mut tree = UbTree::<usize>::new();
193        tree.add_set(BTreeSet::from_iter([0, 1, 2, 3]));
194        tree.add_set(BTreeSet::from_iter([0, 1, 3]));
195        tree.add_set(BTreeSet::from_iter([0, 1, 2]));
196        tree.add_set(BTreeSet::from_iter([2, 3]));
197        assert_eq!(tree.root_nodes().len(), 2);
198        assert!(tree.root_nodes().contains_key(&0));
199        assert!(tree.root_nodes().contains_key(&2));
200
201        // root nodes
202        let e0 = tree.get_node(tree.root_nodes()[&0]);
203        let e2 = tree.get_node(tree.root_nodes()[&2]);
204        assert!(!e0.is_end_of_path());
205        assert_eq!(tree.get_children(e0.children).keys().copied().collect_vec(), vec![1]);
206        assert!(!e2.is_end_of_path());
207        assert_eq!(tree.get_children(e2.children).keys().copied().collect_vec(), vec![3]);
208
209        // first level
210        let e0e1 = tree.get_node(tree.get_children(e0.children)[&1]);
211        let e2e3 = tree.get_node(tree.get_children(e2.children)[&3]);
212        assert!(!e0e1.is_end_of_path());
213        assert_eq!(tree.get_children(e0e1.children).keys().copied().collect_vec(), vec![2, 3]);
214        assert!(e2e3.is_end_of_path());
215        assert_eq!(e2e3.set, Some(BTreeSet::from_iter([2, 3])));
216        assert!(tree.get_children(e2e3.children).is_empty());
217
218        // second level
219        let e0e1e2 = tree.get_node(tree.get_children(e0e1.children)[&2]);
220        assert!(e0e1e2.is_end_of_path());
221        assert_eq!(e0e1e2.set, Some(BTreeSet::from_iter([0, 1, 2])));
222        assert_eq!(tree.get_children(e0e1e2.children).keys().copied().collect_vec(), vec![3]);
223        let e0e1e3 = tree.get_node(tree.get_children(e0e1.children)[&3]);
224        assert!(e0e1e3.is_end_of_path());
225        assert_eq!(e0e1e3.set, Some(BTreeSet::from_iter([0, 1, 3])));
226        assert!(tree.get_children(e0e1e3.children).is_empty());
227
228        // third level
229        let e0e1e2e3 = tree.get_node(tree.get_children(e0e1e2.children)[&3]);
230        assert!(e0e1e2e3.is_end_of_path());
231        assert_eq!(e0e1e2e3.set, Some(BTreeSet::from_iter([0, 1, 2, 3])));
232        assert!(tree.get_children(e0e1e2e3.children).is_empty());
233    }
234
235    #[allow(clippy::cognitive_complexity)]
236    #[test]
237    fn test_contains_subset() {
238        let mut tree = UbTree::<usize>::new();
239        let e0123 = BTreeSet::from_iter([0, 1, 2, 3]);
240        let e013 = BTreeSet::from_iter([0, 1, 3]);
241        let e012 = BTreeSet::from_iter([0, 1, 2]);
242        let e23 = BTreeSet::from_iter([2, 3]);
243        tree.add_set(e0123.clone());
244        tree.add_set(e013.clone());
245        tree.add_set(e012.clone());
246        tree.add_set(e23.clone());
247        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0])), None);
248        assert_eq!(tree.first_subset(&BTreeSet::from_iter([1])), None);
249        assert_eq!(tree.first_subset(&BTreeSet::from_iter([2])), None);
250        assert_eq!(tree.first_subset(&BTreeSet::from_iter([3])), None);
251        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 1])), None);
252        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 2])), None);
253        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 3])), None);
254        assert_eq!(tree.first_subset(&BTreeSet::from_iter([1, 2])), None);
255        assert_eq!(tree.first_subset(&BTreeSet::from_iter([2, 3])), Some(&e23));
256        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 1, 2])), Some(&e012));
257        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 1, 3])), Some(&e013));
258        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 2, 3])), Some(&e23));
259        assert_eq!(tree.first_subset(&BTreeSet::from_iter([1, 2, 3])), Some(&e23));
260        let res0123 = tree.first_subset(&BTreeSet::from_iter([0, 1, 2, 3]));
261        assert!([Some(&e0123), Some(&e013), Some(&e012), Some(&e23)].contains(&res0123));
262        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 4])), None);
263        assert_eq!(tree.first_subset(&BTreeSet::from_iter([1, 4])), None);
264        assert_eq!(tree.first_subset(&BTreeSet::from_iter([2, 4])), None);
265        assert_eq!(tree.first_subset(&BTreeSet::from_iter([3, 4])), None);
266        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 1, 4])), None);
267        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 2, 4])), None);
268        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 3, 4])), None);
269        assert_eq!(tree.first_subset(&BTreeSet::from_iter([1, 2, 4])), None);
270        assert_eq!(tree.first_subset(&BTreeSet::from_iter([1, 2, 4])), None);
271        assert_eq!(tree.first_subset(&BTreeSet::from_iter([2, 3, 4])), Some(&e23));
272        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 1, 2, 4])), Some(&e012));
273        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 1, 3, 4])), Some(&e013));
274        assert_eq!(tree.first_subset(&BTreeSet::from_iter([0, 2, 3, 4])), Some(&e23));
275        assert_eq!(tree.first_subset(&BTreeSet::from_iter([1, 2, 3, 4])), Some(&e23));
276        let res01234 = tree.first_subset(&BTreeSet::from_iter([0, 1, 2, 3, 4]));
277        assert!([Some(&e0123), Some(&e013), Some(&e012), Some(&e23)].contains(&res01234));
278    }
279
280    #[test]
281    fn test_all_sets() {
282        let mut tree = UbTree::<usize>::new();
283        let e0123 = BTreeSet::from_iter([0, 1, 2, 3]);
284        let e013 = BTreeSet::from_iter([0, 1, 3]);
285        let e012 = BTreeSet::from_iter([0, 1, 2]);
286        let e23 = BTreeSet::from_iter([2, 3]);
287        tree.add_set(e0123.clone());
288        assert_eq!(tree.all_sets(), BTreeSet::from_iter([&e0123]));
289        tree.add_set(e013.clone());
290        assert_eq!(tree.all_sets(), BTreeSet::from_iter([&e0123, &e013]));
291        tree.add_set(e012.clone());
292        assert_eq!(tree.all_sets(), BTreeSet::from_iter([&e0123, &e013, &e012]));
293        tree.add_set(e23.clone());
294        assert_eq!(tree.all_sets(), BTreeSet::from_iter([&e0123, &e013, &e012, &e23]));
295    }
296}
logicng/datastructures/ubtree.rs

logicng/datastructures/
ubtree.rs
