indexed-vector 0.1.0

A Rust library that implements vector container that can be indexed with specified function.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

use std::{collections::BTreeMap, ops::RangeBounds, slice::Iter};

use crate::IndexedVector;

/// A simple implementation of `IndexedVector` using `BTreeMap`.
/// It also supports range search.
pub struct BTreeIndexedVector<K, V> {
    map: BTreeMap<K, Vec<V>>,
    key_func: Box<dyn Fn(&V) -> K>,
}

impl<K: Ord, V> BTreeIndexedVector<K, V> {
    /// Create a new `BTreeIndexedVector` from a vector of items.
    /// The `key_func` is used to extract the key from an item.
    pub fn new<F: Fn(&V) -> K + 'static, C: IntoIterator<Item = V>>(data: C, key_func: F) -> Self {
        let mut map = BTreeMap::new();
        for item in data {
            let key = key_func(&item);
            map.entry(key).or_insert_with(Vec::new).push(item);
        }
        Self {
            map,
            key_func: Box::new(key_func),
        }
    }

    /// Search for items in the given range.
    pub fn search_range<R: RangeBounds<K>>(&self, range: R) -> impl Iterator<Item = &V> {
        self.map.range(range).flat_map(|(_, v)| v.iter())
    }
}

impl<K: Ord, V> IndexedVector<K, V> for BTreeIndexedVector<K, V> {
    fn search(&self, key: &K) -> Iter<'_, V> {
        self.map.get(key).map_or([].iter(), |v| v.iter())
    }

    fn insert(&mut self, item: V) {
        let key = (self.key_func)(&item);
        self.map.entry(key).or_insert_with(Vec::new).push(item);
    }
}

#[cfg(test)]
mod test {
    use super::{BTreeIndexedVector, IndexedVector};

    #[test]
    fn test_btree_indexed_vector() {
        let mut map = BTreeIndexedVector::new(
            vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
            Box::new(|x: &i32| x % 3),
        );
        assert_eq!(map.search(&0).collect::<Vec<_>>(), vec![&3, &6, &9]);
        assert_eq!(map.search(&1).collect::<Vec<_>>(), vec![&1, &4, &7, &10]);
        assert_eq!(map.search(&2).collect::<Vec<_>>(), vec![&2, &5, &8]);

        map.insert(11);
        assert_eq!(map.search(&0).collect::<Vec<_>>(), vec![&3, &6, &9]);
        assert_eq!(map.search(&1).collect::<Vec<_>>(), vec![&1, &4, &7, &10]);
        assert_eq!(map.search(&2).collect::<Vec<_>>(), vec![&2, &5, &8, &11]);
    }

    #[test]
    fn test_search_range() {
        let map = BTreeIndexedVector::new(
            vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
            Box::new(|x: &i32| x % 3),
        );
        let mut res = map.search_range(0..2).collect::<Vec<_>>();
        res.sort();
        assert_eq!(res, vec![&1, &3, &4, &6, &7, &9, &10]);
    }
}