algebraeon_sets/structure/
orderings.rs

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