haar_lib/ds/
bitset.rs

1//! 任意サイズのbit列を扱う。
2use std::fmt::Display;
3use std::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign};
4
5type B = u64;
6const B_SIZE: usize = 64;
7
8/// 任意サイズのbit列を扱う。
9#[derive(Clone, Debug)]
10pub struct Bitset {
11    pub(crate) data: Vec<B>,
12    size: usize,
13}
14
15impl Bitset {
16    /// `Bitset`内部で扱う型のBit数
17    pub const B_SIZE: usize = B_SIZE;
18
19    /// 長さ`n`の`Bitset`を構築する。
20    pub fn new(n: usize) -> Self {
21        Self {
22            data: vec![0; n.div_ceil(B_SIZE)],
23            size: n,
24        }
25    }
26
27    /// `n`番目のbitを`val`で設定する。
28    pub fn set(&mut self, n: usize, val: bool) {
29        assert!(n < self.size);
30        if val {
31            unsafe {
32                *self.data.get_unchecked_mut(n / B_SIZE) |= 1 << (n % B_SIZE);
33            }
34        } else {
35            unsafe {
36                *self.data.get_unchecked_mut(n / B_SIZE) &= !(1 << (n % B_SIZE));
37            }
38        }
39    }
40
41    /// `n`番目のbitが`1`ならば`true`を返す。
42    pub fn test(&self, n: usize) -> bool {
43        assert!(n < self.size);
44        unsafe { (self.data.get_unchecked(n / B_SIZE) >> (n % B_SIZE)) & 1 == 1 }
45    }
46
47    /// `n`番目のbitを反転させる。
48    pub fn flip(&mut self, n: usize) {
49        assert!(n < self.size);
50        unsafe {
51            *self.data.get_unchecked_mut(n / B_SIZE) ^= 1 << (n % B_SIZE);
52        }
53    }
54
55    /// `1`が設定されているbitの個数を数える。
56    ///
57    /// **Time complexity** $O(n)$
58    pub fn count_ones(&self) -> u32 {
59        self.data.iter().map(|a| a.count_ones()).sum()
60    }
61
62    /// `0`が設定されているbitの個数を数える。
63    ///
64    /// **Time complexity** $O(n)$
65    pub fn count_zeros(&self) -> u32 {
66        self.size as u32 - self.count_ones()
67    }
68
69    #[allow(missing_docs)]
70    pub fn and_count_ones(&self, rhs: &Self) -> u32 {
71        self.data
72            .iter()
73            .zip(rhs.data.iter())
74            .map(|(a, b)| (a & b).count_ones())
75            .sum()
76    }
77
78    #[allow(missing_docs)]
79    pub fn same_size_xor_assign(&mut self, rhs: &Self) {
80        assert_eq!(self.size, rhs.size);
81        for (a, b) in self.data.iter_mut().zip(rhs.data.iter()) {
82            *a ^= b;
83        }
84    }
85
86    /// bit列の長さを返す。
87    ///
88    /// **Time complexity** $O(1)$
89    pub fn len(&self) -> usize {
90        self.size
91    }
92
93    /// bit列が空ならば`true`を返す。
94    pub fn is_empty(&self) -> bool {
95        self.size == 0
96    }
97
98    /// 末尾に値`val`を追加する。
99    pub fn push(&mut self, val: bool) {
100        if self.size % B_SIZE == 0 {
101            self.data.push(0);
102        }
103        self.size += 1;
104        self.set(self.size - 1, val);
105    }
106
107    /// 末尾の値を削除して返す。
108    pub fn pop(&mut self) -> Option<bool> {
109        if self.size == 0 {
110            None
111        } else {
112            let ret = self.test(self.size - 1);
113            self.set(self.size - 1, false);
114
115            self.size -= 1;
116            if self.size % B_SIZE == 0 {
117                self.data.pop();
118            }
119            Some(ret)
120        }
121    }
122}
123
124impl From<Vec<bool>> for Bitset {
125    fn from(value: Vec<bool>) -> Self {
126        let mut ret = Self::new(value.len());
127
128        for (i, x) in value.chunks(B_SIZE).enumerate() {
129            let mut a = 0;
130
131            for (j, y) in x.iter().enumerate() {
132                if *y {
133                    a |= 1 << j;
134                }
135            }
136
137            ret.data[i] = a;
138        }
139
140        ret
141    }
142}
143
144impl BitAnd for Bitset {
145    type Output = Self;
146
147    fn bitand(mut self, rhs: Self) -> Self::Output {
148        self &= rhs;
149        self
150    }
151}
152
153impl BitAndAssign for Bitset {
154    fn bitand_assign(&mut self, mut rhs: Self) {
155        if self.size > rhs.size {
156            std::mem::swap(self, &mut rhs);
157        }
158
159        for (a, b) in self.data.iter_mut().zip(rhs.data.into_iter()) {
160            *a &= b;
161        }
162    }
163}
164
165impl BitOr for Bitset {
166    type Output = Self;
167
168    fn bitor(mut self, rhs: Self) -> Self::Output {
169        self |= rhs;
170        self
171    }
172}
173
174impl BitOrAssign for Bitset {
175    fn bitor_assign(&mut self, mut rhs: Self) {
176        if self.size < rhs.size {
177            std::mem::swap(self, &mut rhs);
178        }
179
180        for (a, b) in self.data.iter_mut().zip(rhs.data.into_iter()) {
181            *a |= b;
182        }
183    }
184}
185
186impl BitXor for Bitset {
187    type Output = Self;
188
189    fn bitxor(mut self, rhs: Self) -> Self::Output {
190        self ^= rhs;
191        self
192    }
193}
194
195impl BitXorAssign for Bitset {
196    fn bitxor_assign(&mut self, mut rhs: Self) {
197        if self.size < rhs.size {
198            std::mem::swap(self, &mut rhs);
199        }
200
201        for (a, b) in self.data.iter_mut().zip(rhs.data.into_iter()) {
202            *a ^= b;
203        }
204    }
205}
206
207impl Display for Bitset {
208    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209        if let Some(a) = self.data.last() {
210            let w = self.len() % B_SIZE;
211            let w = if w == 0 { B_SIZE } else { w };
212            f.write_fmt(format_args!("{a:0w$b}"))?
213        }
214        for a in self.data.iter().rev().skip(1) {
215            f.write_fmt(format_args!("{a:0B_SIZE$b}"))?
216        }
217        Ok(())
218    }
219}
haar_lib/ds/bitset.rs

haar_lib/ds/
bitset.rs
