cp_library_rs/data_structure/
multiset.rs

1//! 多重集合(Setによる実装)
2
3use rustc_hash::FxHashMap;
4use std::{collections::BTreeSet, fmt::Debug, hash::Hash};
5
6pub struct MultiSet<T> {
7    pub counter: FxHashMap<T, usize>,
8    pub items: BTreeSet<(T, usize)>,
9}
10
11impl<T> Default for MultiSet<T>
12where
13    T: Ord + Hash + Copy,
14{
15    fn default() -> Self {
16        Self::new()
17    }
18}
19
20impl<T> MultiSet<T>
21where
22    T: Ord + Hash + Copy,
23{
24    /// MultiSetを初期化する
25    pub fn new() -> Self {
26        MultiSet {
27            counter: FxHashMap::default(),
28            items: BTreeSet::new(),
29        }
30    }
31
32    /// 要素`x`を追加する
33    pub fn insert(&mut self, x: T) {
34        // カウンターに追加
35        let cnt = self.counter.entry(x).or_insert(0);
36        // setに追加
37        self.items.insert((x, *cnt));
38        // カウント
39        *cnt += 1;
40    }
41
42    /// 要素`x`を削除する
43    pub fn remove(&mut self, x: &T) -> bool {
44        if let Some(v) = self.counter.get_mut(x) {
45            // カウンターをデクリメント
46            if *v > 0 {
47                *v -= 1;
48                // setから削除
49                self.items.remove(&(*x, *v));
50                return true;
51            }
52        }
53        false
54    }
55
56    /// 要素`x`が存在するか判定する
57    pub fn contains(&self, x: &T) -> bool {
58        self.counter.get(x).is_some_and(|cnt| *cnt > 0)
59    }
60
61    /// 先頭の要素を取得する
62    pub fn first(&self) -> Option<&T> {
63        self.items.first().map(|(ref x, _)| x)
64    }
65
66    /// 末尾の要素を取得する
67    pub fn last(&self) -> Option<&T> {
68        self.items.last().map(|(ref x, _)| x)
69    }
70
71    /// `x`の個数をカウントする
72    pub fn count(&self, x: &T) -> usize {
73        match self.counter.get(x) {
74            Some(&v) => v,
75            None => 0,
76        }
77    }
78
79    /// 要素をすべて削除する
80    pub fn clear(&mut self) {
81        self.counter.clear();
82        self.items.clear();
83    }
84
85    pub fn len(&self) -> usize {
86        self.items.len()
87    }
88
89    pub fn is_empty(&self) -> bool {
90        self.len() == 0
91    }
92}
93
94impl<T> MultiSet<T> {
95    pub fn iter(&self) -> impl Iterator<Item = &T> {
96        self.items.iter().map(|(ref x, _)| x)
97    }
98}
99
100impl<T: Ord + Hash + Copy> FromIterator<T> for MultiSet<T> {
101    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
102        let mut multiset = MultiSet::new();
103        for x in iter {
104            multiset.insert(x);
105        }
106        multiset
107    }
108}
109
110impl<T: Debug> Debug for MultiSet<T> {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        f.debug_set().entries(self.iter()).finish()
113    }
114}