indexed_vector/
btree.rs

1use std::{collections::BTreeMap, ops::RangeBounds, slice::Iter};
2
3use crate::IndexedVector;
4
5/// A simple implementation of `IndexedVector` using `BTreeMap`.
6/// It also supports range search.
7pub struct BTreeIndexedVector<K, V> {
8    map: BTreeMap<K, Vec<V>>,
9    key_func: Box<dyn Fn(&V) -> K>,
10}
11
12impl<K: Ord, V> BTreeIndexedVector<K, V> {
13    /// Create a new `BTreeIndexedVector` from a vector of items.
14    /// The `key_func` is used to extract the key from an item.
15    pub fn new<F: Fn(&V) -> K + 'static, C: IntoIterator<Item = V>>(data: C, key_func: F) -> Self {
16        let mut map = BTreeMap::new();
17        for item in data {
18            let key = key_func(&item);
19            map.entry(key).or_insert_with(Vec::new).push(item);
20        }
21        Self {
22            map,
23            key_func: Box::new(key_func),
24        }
25    }
26
27    /// Search for items in the given range.
28    pub fn search_range<R: RangeBounds<K>>(&self, range: R) -> impl Iterator<Item = &V> {
29        self.map.range(range).flat_map(|(_, v)| v.iter())
30    }
31}
32
33impl<K: Ord, V> IndexedVector<K, V> for BTreeIndexedVector<K, V> {
34    fn search(&self, key: &K) -> Iter<'_, V> {
35        self.map.get(key).map_or([].iter(), |v| v.iter())
36    }
37
38    fn insert(&mut self, item: V) {
39        let key = (self.key_func)(&item);
40        self.map.entry(key).or_insert_with(Vec::new).push(item);
41    }
42}
43
44#[cfg(test)]
45mod test {
46    use super::{BTreeIndexedVector, IndexedVector};
47
48    #[test]
49    fn test_btree_indexed_vector() {
50        let mut map = BTreeIndexedVector::new(
51            vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
52            Box::new(|x: &i32| x % 3),
53        );
54        assert_eq!(map.search(&0).collect::<Vec<_>>(), vec![&3, &6, &9]);
55        assert_eq!(map.search(&1).collect::<Vec<_>>(), vec![&1, &4, &7, &10]);
56        assert_eq!(map.search(&2).collect::<Vec<_>>(), vec![&2, &5, &8]);
57
58        map.insert(11);
59        assert_eq!(map.search(&0).collect::<Vec<_>>(), vec![&3, &6, &9]);
60        assert_eq!(map.search(&1).collect::<Vec<_>>(), vec![&1, &4, &7, &10]);
61        assert_eq!(map.search(&2).collect::<Vec<_>>(), vec![&2, &5, &8, &11]);
62    }
63
64    #[test]
65    fn test_search_range() {
66        let map = BTreeIndexedVector::new(
67            vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
68            Box::new(|x: &i32| x % 3),
69        );
70        let mut res = map.search_range(0..2).collect::<Vec<_>>();
71        res.sort();
72        assert_eq!(res, vec![&1, &3, &4, &6, &7, &9, &10]);
73    }
74}
indexed_vector/btree.rs

indexed_vector/
btree.rs
