sorted_groups/
lib.rs

1//! `sorted-groups` implement a data structure to store elements in sorted groups while maintaining the order of elements in each group.
2//!
3//! # Usage
4//!
5//! First, add the `sorted_groups` crate as a dependency:
6//! ```sh
7//! cargo add sorted_groups
8//! ```
9//!
10//! ```
11//! use sorted_groups::SortedGroups;
12//!
13//! #[derive(PartialEq, Eq, Ord, Debug)]
14//! struct Element {
15//!    group: i32,
16//!    value: i32,
17//! }
18//!
19//! impl PartialOrd for Element {
20//!    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
21//!        Some(self.cmp(other))
22//!    }
23//! }
24//!
25//! // Elements will be grouped by the `group` field
26//! let sorted_groups = SortedGroups::<i32, Element>::new(vec![
27//!    Element { group: 1, value: 1 },
28//!    Element { group: 1, value: 2 },
29//!    Element { group: 2, value: 1 },
30//! ], |e| e.group);
31//!
32//! // `len` returns the total number of elements
33//! assert_eq!(sorted_groups.len(), 3);
34//! // `groups_len` returns the number of groups
35//! assert_eq!(sorted_groups.groups_len(), 2);
36//! // `iter` returns an iterator over groups and elements
37//! let mut iter = sorted_groups.iter();
38//! assert_eq!(iter.next(), Some((&1, &Element { group: 1, value: 1 })));
39//! assert_eq!(iter.next(), Some((&1, &Element { group: 1, value: 2 })));
40//! assert_eq!(iter.next(), Some((&2, &Element { group: 2, value: 1 })));
41//! assert_eq!(iter.next(), None);
42//! ```
43//!
44use std::collections::{btree_map::BTreeMap, btree_set, BTreeSet};
45
46#[derive(Clone, Debug)]
47pub struct SortedGroups<G, E>
48where
49    G: Ord,
50    E: Ord,
51{
52    groups: BTreeMap<G, BTreeSet<E>>,
53}
54
55impl<G, E> SortedGroups<G, E>
56where
57    G: Ord,
58    E: Ord,
59{
60    pub fn new(
61        elements: impl IntoIterator<Item = E>,
62        group_from_element: impl Fn(&E) -> G,
63    ) -> Self {
64        let mut groups = BTreeMap::<G, BTreeSet<E>>::new();
65        for element in elements {
66            groups
67                .entry(group_from_element(&element))
68                .or_default()
69                .insert(element);
70        }
71        Self { groups }
72    }
73
74    pub fn len(&self) -> usize {
75        self.groups.values().map(|v| v.len()).sum()
76    }
77
78    pub fn is_empty(&self) -> bool {
79        self.len() == 0
80    }
81
82    pub fn get(&self, index: usize) -> Option<(&G, &E)> {
83        self.iter().nth(index)
84    }
85
86    pub fn groups_len(&self) -> usize {
87        self.groups.len()
88    }
89
90    pub fn iter_groups(&self) -> impl Iterator<Item = (&G, &BTreeSet<E>)> {
91        self.groups.iter()
92    }
93}
94
95pub struct SortedGroupsIter<'a, G, E> {
96    // Iterator over groups
97    groups_iter: std::collections::btree_map::Iter<'a, G, BTreeSet<E>>,
98    // Current group and its iterator
99    current_group: Option<(&'a G, btree_set::Iter<'a, E>)>,
100}
101
102impl<G, E> SortedGroups<G, E>
103where
104    G: Ord,
105    E: Ord,
106{
107    pub fn iter(&self) -> SortedGroupsIter<'_, G, E> {
108        let mut groups_iter = self.groups.iter();
109        let current_group = groups_iter.next().map(|(g, v)| (g, v.iter()));
110
111        SortedGroupsIter {
112            groups_iter,
113            current_group,
114        }
115    }
116}
117
118impl<'a, G, E> Iterator for SortedGroupsIter<'a, G, E>
119where
120    G: Ord,
121    E: Ord,
122{
123    type Item = (&'a G, &'a E);
124
125    fn next(&mut self) -> Option<Self::Item> {
126        loop {
127            match &mut self.current_group {
128                Some((group, iter)) => {
129                    if let Some(element) = iter.next() {
130                        return Some((*group, element));
131                    } else {
132                        // Current group is exhausted, move to next group
133                        self.current_group = self.groups_iter.next().map(|(g, v)| (g, v.iter()));
134                    }
135                }
136                None => return None,
137            }
138        }
139    }
140}
141
142// Implement IntoIterator for reference
143impl<'a, G, E> IntoIterator for &'a SortedGroups<G, E>
144where
145    G: Ord,
146    E: Ord,
147{
148    type Item = (&'a G, &'a E);
149    type IntoIter = SortedGroupsIter<'a, G, E>;
150
151    fn into_iter(self) -> Self::IntoIter {
152        self.iter()
153    }
154}
155
156impl<G, E> PartialEq for SortedGroups<G, E>
157where
158    G: Ord,
159    E: Ord,
160{
161    fn eq(&self, other: &Self) -> bool {
162        self.groups.eq(&other.groups)
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    #[derive(PartialEq, Eq, Ord, Debug)]
171    struct Element {
172        group: i32,
173        value: i32,
174    }
175
176    impl PartialOrd for Element {
177        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
178            Some(self.cmp(other))
179        }
180    }
181
182    #[test]
183    fn test_empty_sorted_groups() {
184        let sorted_groups = SortedGroups::<i32, Element>::new(vec![].into_iter(), |e| e.group);
185        assert_eq!(sorted_groups.len(), 0);
186    }
187
188    #[test]
189    fn test_insert_sorted_groups() {
190        let sorted_groups = SortedGroups::<i32, Element>::new(
191            vec![
192                Element { group: 1, value: 1 },
193                Element { group: 1, value: 2 },
194                Element { group: 2, value: 1 },
195            ],
196            |e| e.group,
197        );
198
199        assert_eq!(sorted_groups.len(), 3);
200        assert_eq!(sorted_groups.groups_len(), 2);
201        let mut iter = sorted_groups.iter();
202        assert_eq!(iter.next(), Some((&1, &Element { group: 1, value: 1 })));
203        assert_eq!(iter.next(), Some((&1, &Element { group: 1, value: 2 })));
204        assert_eq!(iter.next(), Some((&2, &Element { group: 2, value: 1 })));
205        assert_eq!(iter.next(), None);
206    }
207}
sorted_groups/lib.rs

sorted_groups/
lib.rs
