haar_lib/ds/
multiset.rs

1//! 同一要素を複数個挿入可能な`Set`
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/abc308/tasks/abc308_f>
5use std::collections::BTreeMap;
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}