haar_lib/ds/
usize_set.rs

1//! `usize`を用いた集合表現
2//!
3//! # Problems
4//! - <https://atcoder.jp/contests/abc142/tasks/abc142_e>
5use std::ops::{BitAnd, BitOr, BitXor, Sub};
6
7/// `usize`のビット数個の要素をもつ集合を表す。
8#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub struct UsizeSet(pub usize);
10
11impl UsizeSet {
12    /// `i`を追加した集合を返す。
13    #[inline]
14    pub fn set(self, i: usize) -> Self {
15        Self(self.0 | (1 << i))
16    }
17
18    /// `i`を削除した集合を返す。
19    #[inline]
20    pub fn reset(self, i: usize) -> Self {
21        Self(self.0 & !(1 << i))
22    }
23
24    /// `i`が集合に含まれていなければ、`i`を追加した集合を返す。
25    ///
26    /// `i`が集合に含まれていれば、`i`を削除した集合を返す。
27    #[inline]
28    pub fn flip(self, i: usize) -> Self {
29        Self(self.0 ^ (1 << i))
30    }
31
32    /// `i`が集合に含まれているかを判定する。
33    #[inline]
34    pub fn contains(self, i: usize) -> bool {
35        (self.0 >> i) & 1 == 1
36    }
37
38    /// `0`から`n`までを要素に含む集合を得る。
39    #[inline]
40    pub fn fill(n: usize) -> Self {
41        assert!(n <= usize::BITS as usize);
42        if n == usize::BITS as usize {
43            Self(!0)
44        } else {
45            Self(!(!0 << n))
46        }
47    }
48
49    /// `self`から`rhs`を引いた差集合を返す。
50    #[inline]
51    pub fn difference(self, rhs: Self) -> Self {
52        Self(self.0 & !rhs.0)
53    }
54
55    /// 2つの集合の和集合を返す。
56    #[inline]
57    pub fn union(self, rhs: Self) -> Self {
58        Self(self.0 | rhs.0)
59    }
60
61    /// 2つの集合の共通部分を返す。
62    #[inline]
63    pub fn intersection(self, rhs: Self) -> Self {
64        Self(self.0 & rhs.0)
65    }
66
67    /// 2つの集合の対象差を返す。
68    #[inline]
69    pub fn symmetric_difference(self, rhs: Self) -> Self {
70        Self(self.0 ^ rhs.0)
71    }
72
73    /// 集合が空かを判定する。
74    #[inline]
75    pub fn is_empty(self) -> bool {
76        self.0 == 0
77    }
78
79    /// 集合の大きさを返す。
80    #[inline]
81    pub fn len(self) -> usize {
82        self.0.count_ones() as usize
83    }
84}
85
86impl BitAnd for UsizeSet {
87    type Output = Self;
88    fn bitand(self, rhs: Self) -> Self::Output {
89        self.intersection(rhs)
90    }
91}
92
93impl BitOr for UsizeSet {
94    type Output = Self;
95    fn bitor(self, rhs: Self) -> Self::Output {
96        self.union(rhs)
97    }
98}
99
100impl BitXor for UsizeSet {
101    type Output = Self;
102    fn bitxor(self, rhs: Self) -> Self::Output {
103        self.symmetric_difference(rhs)
104    }
105}
106
107impl Sub for UsizeSet {
108    type Output = Self;
109    fn sub(self, rhs: Self) -> Self::Output {
110        self.difference(rhs)
111    }
112}
113
114impl From<Vec<usize>> for UsizeSet {
115    fn from(value: Vec<usize>) -> Self {
116        let mut ret = Self(0);
117        for a in value {
118            ret = ret.set(a);
119        }
120        ret
121    }
122}
123
124impl From<UsizeSet> for Vec<usize> {
125    fn from(value: UsizeSet) -> Self {
126        (0..usize::BITS as usize)
127            .filter(|i| value.contains(*i))
128            .collect()
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use std::collections::BTreeSet;
135    use std::iter::FromIterator;
136
137    use super::*;
138    use rand::{seq::SliceRandom, Rng};
139
140    #[test]
141    fn test() {
142        let mut rng = rand::thread_rng();
143
144        #[allow(non_snake_case)]
145        let U: Vec<usize> = (0..usize::BITS as usize).collect();
146
147        for _ in 0..100 {
148            let count = rng.gen::<usize>() % 65;
149            let a: Vec<_> = U.choose_multiple(&mut rng, count).cloned().collect();
150            let b: Vec<_> = U.choose_multiple(&mut rng, count).cloned().collect();
151
152            let a_ans = BTreeSet::from_iter(a.clone());
153            let b_ans = BTreeSet::from_iter(b.clone());
154
155            let a_res = UsizeSet::from(a);
156            let b_res = UsizeSet::from(b);
157
158            let c_ans = &a_ans & &b_ans;
159            let c_res = a_res & b_res;
160            assert_eq!(BTreeSet::from_iter(Vec::from(c_res)), c_ans);
161
162            let c_ans = &a_ans | &b_ans;
163            let c_res = a_res | b_res;
164            assert_eq!(BTreeSet::from_iter(Vec::from(c_res)), c_ans);
165
166            let c_ans = &a_ans ^ &b_ans;
167            let c_res = a_res ^ b_res;
168            assert_eq!(BTreeSet::from_iter(Vec::from(c_res)), c_ans);
169
170            let c_ans = &a_ans - &b_ans;
171            let c_res = a_res - b_res;
172            assert_eq!(BTreeSet::from_iter(Vec::from(c_res)), c_ans);
173        }
174    }
175}