algebraeon_sets/structure/
orderings.rs

1use super::EqSignature;
2use algebraeon_macros::{signature_meta_trait, skip_meta};
3use std::{borrow::Borrow, cmp::Ordering};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum MergedSource {
7    First,
8    Second,
9}
10
11struct VecMerger<'s, X, O: OrdSignature, K: Fn(&X) -> &O::Set> {
12    ordering: &'s O,
13    i: usize,
14    a: Vec<Option<X>>,
15    j: usize,
16    b: Vec<Option<X>>,
17    key: K,
18}
19
20impl<'s, X, O: OrdSignature + 's, K: Fn(&X) -> &O::Set> VecMerger<'s, X, O, K> {
21    fn new(ordering: &'s O, a: Vec<X>, b: Vec<X>, key: K) -> Self {
22        Self {
23            ordering,
24            i: 0,
25            a: a.into_iter().map(|x| Some(x)).collect(),
26            j: 0,
27            b: b.into_iter().map(|x| Some(x)).collect(),
28            key,
29        }
30    }
31}
32
33impl<'s, X, O: OrdSignature + 's, K: Fn(&X) -> &O::Set> Iterator for VecMerger<'s, X, O, K> {
34    type Item = (MergedSource, X);
35
36    fn next(&mut self) -> Option<Self::Item> {
37        match (self.i < self.a.len(), self.j < self.b.len()) {
38            (true, true) => {
39                match self.ordering.cmp(
40                    (self.key)(self.a[self.i].as_ref().unwrap()).borrow(),
41                    (self.key)(self.b[self.j].as_ref().unwrap()).borrow(),
42                ) {
43                    Ordering::Less | Ordering::Equal => {
44                        let a_item = self.a[self.i].take().unwrap();
45                        self.i += 1;
46                        Some((MergedSource::First, a_item))
47                    }
48                    Ordering::Greater => {
49                        let b_item = self.b[self.j].take().unwrap();
50                        self.j += 1;
51                        Some((MergedSource::Second, b_item))
52                    }
53                }
54            }
55            (true, false) => {
56                let a_item = self.a[self.i].take().unwrap();
57                self.i += 1;
58                Some((MergedSource::First, a_item))
59            }
60            (false, true) => {
61                let b_item = self.b[self.j].take().unwrap();
62                self.j += 1;
63                Some((MergedSource::Second, b_item))
64            }
65            (false, false) => None,
66        }
67    }
68}
69
70use super::MetaType;
71
72#[signature_meta_trait]
73pub trait PartialOrdSignature: EqSignature {
74    fn partial_cmp(&self, a: &Self::Set, b: &Self::Set) -> Option<Ordering>;
75}
76
77#[signature_meta_trait]
78pub trait OrdSignature: PartialOrdSignature {
79    fn cmp(&self, a: &Self::Set, b: &Self::Set) -> Ordering;
80
81    fn max(&self, a: &Self::Set, b: &Self::Set) -> Self::Set {
82        let c = self.cmp(a, b);
83        match c {
84            Ordering::Less | Ordering::Equal => b.clone(),
85            Ordering::Greater => a.clone(),
86        }
87    }
88
89    fn min(&self, a: &Self::Set, b: &Self::Set) -> Self::Set {
90        let c = self.cmp(a, b);
91        match c {
92            Ordering::Less | Ordering::Equal => a.clone(),
93            Ordering::Greater => b.clone(),
94        }
95    }
96
97    fn is_sorted(&self, a: &[impl Borrow<Self::Set>]) -> bool {
98        for i in 1..a.len() {
99            if self.cmp(a[i - 1].borrow(), a[i].borrow()) == Ordering::Greater {
100                return false;
101            }
102        }
103        true
104    }
105
106    fn is_sorted_by_key<X>(&self, a: &[X], key: impl Fn(&X) -> &Self::Set) -> bool {
107        self.is_sorted(&a.iter().map(key).collect::<Vec<_>>())
108    }
109
110    fn is_sorted_and_unique(&self, a: &[impl Borrow<Self::Set>]) -> bool {
111        for i in 1..a.len() {
112            if self.cmp(a[i - 1].borrow(), a[i].borrow()) != Ordering::Less {
113                return false;
114            }
115        }
116        true
117    }
118
119    fn is_sorted_and_unique_by_key<X>(&self, a: &[X], key: impl Fn(&X) -> &Self::Set) -> bool {
120        self.is_sorted_and_unique(&a.iter().map(key).collect::<Vec<_>>())
121    }
122
123    fn binary_search(&self, v: &[impl Borrow<Self::Set>], target: &Self::Set) -> bool {
124        self.binary_search_by_key(v, target, |x| x.borrow())
125            .is_some()
126    }
127
128    fn binary_search_by_key<'x, X>(
129        &self,
130        v: &'x [X],
131        target: &Self::Set,
132        key: impl Fn(&X) -> &Self::Set,
133    ) -> Option<&'x X> {
134        debug_assert!(self.is_sorted_by_key(v, &key));
135        if v.is_empty() {
136            return None;
137        }
138        let mut a = 0;
139        let mut b = v.len() - 1;
140
141        if self.equal(target, key(&v[a])) {
142            return Some(&v[a]);
143        }
144
145        if self.equal(target, key(&v[b])) {
146            return Some(&v[b]);
147        }
148
149        while b - a >= 2 {
150            println!("{:?} {:?}", a, b);
151
152            let m = (a + b) / 2;
153            let m_key = key(&v[m]);
154            match self.cmp(target, m_key) {
155                Ordering::Less => b = m,
156                Ordering::Equal => {
157                    return Some(&v[m]);
158                }
159                Ordering::Greater => a = m,
160            }
161        }
162
163        None
164    }
165
166    #[skip_meta]
167    fn merge_sorted<'s, S: Borrow<Self::Set> + 's>(
168        &'s self,
169        a: Vec<S>,
170        b: Vec<S>,
171    ) -> impl Iterator<Item = (MergedSource, S)> {
172        self.merge_sorted_by_key(a, b, |x| (*x).borrow())
173    }
174
175    #[skip_meta]
176    fn merge_sorted_by_key<X>(
177        &self,
178        a: Vec<X>,
179        b: Vec<X>,
180        key: impl Fn(&X) -> &Self::Set,
181    ) -> impl Iterator<Item = (MergedSource, X)> {
182        debug_assert!(self.is_sorted_by_key(&a.iter().collect::<Vec<_>>(), |x| key(x)));
183        debug_assert!(self.is_sorted_by_key(&b.iter().collect::<Vec<_>>(), |x| key(x)));
184        VecMerger::new(self, a, b, key)
185    }
186
187    fn sort<S: Borrow<Self::Set>>(&self, a: Vec<S>) -> Vec<S> {
188        self.sort_by_key(a, &|x| x.borrow())
189    }
190
191    #[skip_meta]
192    fn sort_by_key<X>(&self, mut a: Vec<X>, key: &impl Fn(&X) -> &Self::Set) -> Vec<X> {
193        match a.len() {
194            0 | 1 => a,
195            2 => match self.cmp(key(&a[0]).borrow(), key(&a[1]).borrow()) {
196                Ordering::Less | Ordering::Equal => a,
197                Ordering::Greater => {
198                    a.swap(0, 1);
199                    a
200                }
201            },
202            n => {
203                // merge sort
204                let k = n / 2;
205
206                let a1 = a.split_off(k);
207                let a0 = a;
208
209                let a0_sorted = self.sort_by_key(a0, key);
210                let a1_sorted = self.sort_by_key(a1, key);
211
212                self.merge_sorted_by_key(a0_sorted, a1_sorted, key)
213                    .map(|(_, s)| s)
214                    .collect()
215            }
216        }
217    }
218
219    fn sort_by_cached_by<X: 'static>(&self, a: Vec<X>, key: impl Fn(&X) -> Self::Set) -> Vec<X>
220    where
221        Self::Set: 'static,
222    {
223        let mut a = a.into_iter().map(|x| (x, None)).collect::<Vec<_>>();
224        for i in 0..a.len() {
225            let (x, k) = a.get_mut(i).unwrap();
226            *k = Some(key(x));
227        }
228        self.sort_by_key(a, &|(_, k)| k.as_ref().unwrap())
229            .into_iter()
230            .map(|(x, _)| x)
231            .collect()
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use crate::structure::*;
238
239    use super::*;
240
241    #[derive(Debug, Clone, PartialEq, Eq)]
242    struct UsizeStructure {}
243
244    impl Signature for UsizeStructure {}
245
246    impl SetSignature for UsizeStructure {
247        type Set = usize;
248
249        fn is_element(&self, _x: &Self::Set) -> Result<(), String> {
250            Ok(())
251        }
252    }
253
254    impl EqSignature for UsizeStructure {
255        fn equal(&self, a: &Self::Set, b: &Self::Set) -> bool {
256            a == b
257        }
258    }
259
260    impl PartialOrdSignature for UsizeStructure {
261        fn partial_cmp(&self, a: &Self::Set, b: &Self::Set) -> Option<Ordering> {
262            Some(self.cmp(a, b))
263        }
264    }
265
266    impl OrdSignature for UsizeStructure {
267        fn cmp(&self, a: &Self::Set, b: &Self::Set) -> Ordering {
268            usize::cmp(a, b)
269        }
270    }
271
272    #[test]
273    fn ordering_structure() {
274        let s = UsizeStructure {};
275
276        assert_eq!(s.cmp(&2, &2), Ordering::Equal);
277        assert_eq!(s.cmp(&1, &2), Ordering::Less);
278        assert_eq!(s.cmp(&2, &1), Ordering::Greater);
279
280        assert!(s.is_sorted(&Vec::<usize>::new()));
281        assert!(s.is_sorted(&[7]));
282        assert!(s.is_sorted(&[1, 2]));
283        assert!(!s.is_sorted(&[2, 1]));
284        assert!(s.is_sorted(&[1, 2, 3]));
285        assert!(s.is_sorted(&[1, 1, 1]));
286        assert!(!s.is_sorted(&[1, 2, 1]));
287
288        assert_eq!(
289            s.merge_sorted(Vec::<usize>::new(), Vec::<usize>::new())
290                .collect::<Vec<_>>(),
291            vec![]
292        );
293        {
294            let a = vec![2, 4, 5, 7];
295            let b = vec![1, 3, 6, 7, 8];
296            let c = s.merge_sorted(a, b).collect::<Vec<_>>();
297            assert_eq!(
298                c,
299                vec![
300                    (MergedSource::Second, 1),
301                    (MergedSource::First, 2),
302                    (MergedSource::Second, 3),
303                    (MergedSource::First, 4),
304                    (MergedSource::First, 5),
305                    (MergedSource::Second, 6),
306                    (MergedSource::First, 7),
307                    (MergedSource::Second, 7),
308                    (MergedSource::Second, 8)
309                ]
310            );
311        }
312
313        assert_eq!(s.sort(Vec::<usize>::new()), Vec::<usize>::new());
314        assert_eq!(s.sort(vec![7]), vec![7]);
315        assert_eq!(s.sort(vec![7, 7]), vec![7, 7]);
316        assert_eq!(s.sort(vec![1, 2]), vec![1, 2]);
317        assert_eq!(s.sort(vec![2, 1]), vec![1, 2]);
318        assert_eq!(s.sort(vec![7, 6, 5, 4, 3, 2, 1]), vec![1, 2, 3, 4, 5, 6, 7]);
319        assert_eq!(
320            s.sort(vec![3, 3, 2, 2, 2, 1, 1, 1, 1]),
321            vec![1, 1, 1, 1, 2, 2, 2, 3, 3]
322        );
323        assert_eq!(
324            s.sort_by_cached_by(vec![3, 3, 2, 2, 2, 1, 1, 1, 1], |x| 5 - x),
325            vec![3, 3, 2, 2, 2, 1, 1, 1, 1]
326        );
327
328        let v = vec![1, 2, 3, 3, 3, 4, 4, 5, 7, 8, 9, 10, 11];
329        assert!(!s.binary_search(&v, &0));
330        assert!(s.binary_search(&v, &1));
331        assert!(s.binary_search(&v, &3));
332        assert!(s.binary_search(&v, &4));
333        assert!(s.binary_search(&v, &5));
334        assert!(!s.binary_search(&v, &6));
335        assert!(s.binary_search(&v, &7));
336        assert!(s.binary_search(&v, &11));
337        assert!(!s.binary_search(&v, &12));
338    }
339}
algebraeon_sets/structure/orderings.rs

algebraeon_sets/structure/
orderings.rs
