haar_lib/ds/
multiset.rs

1//! 同一要素を複数個挿入可能な`Set`
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/abc308/tasks/abc308_f>
5use std::{collections::BTreeMap, ops::Bound};
6
7/// 同一要素を複数個挿入可能な`Set`
8#[derive(Debug, Clone, Default)]
9pub struct MultiSet<T> {
10    map: BTreeMap<T, usize>,
11    size: usize,
12}
13
14impl<T: Ord + Eq + Clone> MultiSet<T> {
15    /// [`MultiSet<T>`]を生成する。
16    pub fn new() -> Self {
17        Self {
18            map: BTreeMap::new(),
19            size: 0,
20        }
21    }
22
23    /// 値`value`を挿入する。
24    pub fn insert(&mut self, value: T) {
25        *self.map.entry(value).or_default() += 1;
26        self.size += 1;
27    }
28
29    /// 値`value`を*一つだけ*削除する。
30    pub fn remove(&mut self, value: T) -> bool {
31        if let Some(count) = self.map.get_mut(&value) {
32            *count -= 1;
33            self.size -= 1;
34
35            if *count == 0 {
36                self.map.remove(&value);
37            }
38
39            true
40        } else {
41            false
42        }
43    }
44
45    /// 先頭の要素を返す。
46    pub fn first(&self) -> Option<T> {
47        self.map.iter().next().map(|(k, _)| k.clone())
48    }
49
50    /// 末尾の要素を返す。
51    pub fn last(&self) -> Option<T> {
52        self.map.iter().next_back().map(|(k, _)| k.clone())
53    }
54
55    /// 末尾の要素を*一つだけ*削除して返す。
56    pub fn pop_last(&mut self) -> Option<T> {
57        if let Some((k, v)) = self.map.iter_mut().next_back() {
58            *v -= 1;
59            self.size -= 1;
60
61            let k = k.clone();
62
63            if *v == 0 {
64                self.map.remove(&k);
65            }
66
67            Some(k)
68        } else {
69            None
70        }
71    }
72
73    /// 先頭の要素を*一つだけ*削除して返す。
74    pub fn pop_first(&mut self) -> Option<T> {
75        if let Some((k, v)) = self.map.iter_mut().next() {
76            *v -= 1;
77            self.size -= 1;
78
79            let k = k.clone();
80
81            if *v == 0 {
82                self.map.remove(&k);
83            }
84
85            Some(k)
86        } else {
87            None
88        }
89    }
90
91    /// 値`value`が含まれていれば、`true`を返す。
92    pub fn contains(&self, value: &T) -> bool {
93        self.map.contains_key(value)
94    }
95
96    /// 値`value`が含まれている個数を返す。
97    pub fn count(&self, value: &T) -> usize {
98        self.map.get(value).cloned().unwrap_or(0)
99    }
100
101    /// 要素数を返す。
102    pub fn len(&self) -> usize {
103        self.size
104    }
105
106    /// 要素数が0ならば、`true`を返す。
107    pub fn is_empty(&self) -> bool {
108        self.map.is_empty()
109    }
110
111    /// `value`以上の最小の要素を返す。
112    pub fn ge(&self, value: &T) -> Option<&T> {
113        self.map.range(value..).next().map(|(k, _)| k)
114    }
115
116    /// `value`以下の最大の要素を返す。
117    pub fn le(&self, value: &T) -> Option<&T> {
118        self.map.range(..=value).next_back().map(|(k, _)| k)
119    }
120
121    /// `value`より大きい最小の要素を返す。
122    pub fn gt(&self, value: &T) -> Option<&T> {
123        self.map
124            .range((Bound::Excluded(value), Bound::Unbounded))
125            .next()
126            .map(|(k, _)| k)
127    }
128
129    /// `value`より小さい最大の要素を返す。
130    pub fn lt(&self, value: &T) -> Option<&T> {
131        self.map.range(..value).next_back().map(|(k, _)| k)
132    }
133}
134
135impl<T: Ord + Eq + Clone> FromIterator<T> for MultiSet<T> {
136    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
137        let mut ret = Self::new();
138        for x in iter {
139            ret.insert(x);
140        }
141        ret
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn test() {
151        let set = MultiSet::from_iter([1, 1, 2, 3, 5, 5, 6, 7, 10]);
152
153        assert_eq!(set.ge(&1), Some(&1));
154        assert_eq!(set.le(&1), Some(&1));
155        assert_eq!(set.gt(&1), Some(&2));
156        assert_eq!(set.lt(&1), None);
157
158        assert_eq!(set.ge(&10), Some(&10));
159        assert_eq!(set.le(&10), Some(&10));
160        assert_eq!(set.gt(&10), None);
161        assert_eq!(set.lt(&10), Some(&7));
162
163        assert_eq!(set.ge(&4), Some(&5));
164        assert_eq!(set.le(&4), Some(&3));
165        assert_eq!(set.gt(&4), Some(&5));
166        assert_eq!(set.lt(&4), Some(&3));
167    }
168}
haar_lib/ds/multiset.rs

haar_lib/ds/
multiset.rs
