btreemultimap_value_ord/
lib.rs

1#[cfg(feature = "serde")]
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4
5#[derive(Clone, Debug, Eq, PartialEq, Hash)]
6#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
7pub struct BTreeMultiMap<Key: Ord, Value: Ord> {
8    map: BTreeMap<Key, BTreeMap<Value, usize>>,
9}
10
11impl<Key: Ord, Value: Ord> BTreeMultiMap<Key, Value> {
12    pub fn new() -> Self {
13        Default::default()
14    }
15
16    pub fn clear(&mut self) {
17        self.map.clear()
18    }
19
20    pub fn insert(&mut self, key: Key, value: Value) {
21        if let Some(map) = self.map.get_mut(&key) {
22            if let Some(count) = map.get_mut(&value) {
23                *count += 1;
24            } else {
25                map.insert(value, 1);
26            }
27        } else {
28            self.map.insert(key, BTreeMap::from([(value, 1)]));
29        }
30    }
31
32    pub fn get(&self, key: &Key) -> Option<&BTreeMap<Value, usize>> {
33        self.map.get(key)
34    }
35
36    pub fn get_value_count(&self, key: &Key, value: &Value) -> usize {
37        if let Some(map) = self.map.get(key) {
38            map.get(value).copied().unwrap_or(0)
39        } else {
40            0
41        }
42    }
43
44    pub fn remove(&mut self, key: &Key) -> Option<BTreeMap<Value, usize>> {
45        self.map.remove(key)
46    }
47
48    pub fn remove_key_value(&mut self, key: &Key, value: &Value) -> bool {
49        if let Some(map) = self.map.get_mut(key) {
50            if let Some(count) = map.get_mut(value) {
51                *count -= 1;
52                if *count == 0 {
53                    map.remove(value);
54                }
55                true
56            } else {
57                false
58            }
59        } else {
60            false
61        }
62    }
63}
64
65impl<Key: Ord, Value: Ord> Default for BTreeMultiMap<Key, Value> {
66    fn default() -> Self {
67        Self {
68            map: Default::default(),
69        }
70    }
71}
72
73impl<Key: Ord, Value: Ord> FromIterator<(Key, Value)> for BTreeMultiMap<Key, Value> {
74    fn from_iter<T: IntoIterator<Item = (Key, Value)>>(iter: T) -> Self {
75        let mut result = Self::default();
76        for (key, value) in iter {
77            result.insert(key, value);
78        }
79        result
80    }
81}