competitive/data_structure/
bitset.rs

1#![allow(clippy::suspicious_op_assign_impl)]
2
3use std::ops::{
4    BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not, Shl, ShlAssign, Shr,
5    ShrAssign,
6};
7
8#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub struct BitSet {
10    size: usize,
11    bits: Vec<u64>,
12}
13impl BitSet {
14    pub fn new(size: usize) -> Self {
15        Self {
16            size,
17            bits: vec![0; (size + 63) / 64],
18        }
19    }
20    pub fn ones(size: usize) -> Self {
21        let mut self_ = Self {
22            size,
23            bits: vec![u64::MAX; (size + 63) / 64],
24        };
25        self_.trim();
26        self_
27    }
28    #[inline]
29    pub fn get(&self, i: usize) -> bool {
30        self.bits[i >> 6] & (1 << (i & 63)) != 0
31    }
32    #[inline]
33    pub fn set(&mut self, i: usize, b: bool) {
34        if b {
35            self.bits[i >> 6] |= 1 << (i & 63);
36        } else {
37            self.bits[i >> 6] &= !(1 << (i & 63));
38        }
39    }
40    #[inline]
41    pub fn count_ones(&self) -> u64 {
42        self.bits.iter().map(|x| x.count_ones() as u64).sum()
43    }
44    #[inline]
45    pub fn count_zeros(&self) -> u64 {
46        self.size as u64 - self.count_ones()
47    }
48    #[inline]
49    fn trim(&mut self) {
50        if self.size & 63 != 0 {
51            if let Some(x) = self.bits.last_mut() {
52                *x &= 0xffff_ffff_ffff_ffff >> (64 - (self.size & 63));
53            }
54        }
55    }
56    #[inline]
57    pub fn shl_bitor_assign(&mut self, rhs: usize) {
58        let n = self.bits.len();
59        let k = rhs >> 6;
60        let d = rhs & 63;
61        if k < n {
62            if d == 0 {
63                for i in (0..n - k).rev() {
64                    self.bits[i + k] |= self.bits[i];
65                }
66            } else {
67                for i in (1..n - k).rev() {
68                    self.bits[i + k] |= (self.bits[i] << d) | (self.bits[i - 1] >> (64 - d));
69                }
70                self.bits[k] |= self.bits[0] << d;
71            }
72            self.trim();
73        }
74    }
75    #[inline]
76    pub fn shr_bitor_assign(&mut self, rhs: usize) {
77        let n = self.bits.len();
78        let k = rhs >> 6;
79        let d = rhs & 63;
80        if k < n {
81            if d == 0 {
82                for i in k..n {
83                    self.bits[i - k] |= self.bits[i];
84                }
85            } else {
86                for i in k..n - 1 {
87                    self.bits[i - k] |= (self.bits[i] >> d) | (self.bits[i + 1] << (64 - d));
88                }
89                self.bits[n - k - 1] |= self.bits[n - 1] >> d;
90            }
91        }
92    }
93}
94impl ShlAssign<usize> for BitSet {
95    #[inline]
96    fn shl_assign(&mut self, rhs: usize) {
97        let n = self.bits.len();
98        let k = rhs >> 6;
99        let d = rhs & 63;
100        if k >= n {
101            for x in self.bits.iter_mut() {
102                *x = 0;
103            }
104        } else {
105            if d == 0 {
106                for i in (0..n - k).rev() {
107                    self.bits[i + k] = self.bits[i];
108                }
109            } else {
110                for i in (1..n - k).rev() {
111                    self.bits[i + k] = (self.bits[i] << d) | (self.bits[i - 1] >> (64 - d));
112                }
113                self.bits[k] = self.bits[0] << d;
114            }
115            for x in self.bits[..k].iter_mut() {
116                *x = 0;
117            }
118            self.trim();
119        }
120    }
121}
122impl Shl<usize> for BitSet {
123    type Output = Self;
124    #[inline]
125    fn shl(mut self, rhs: usize) -> Self::Output {
126        self <<= rhs;
127        self
128    }
129}
130impl ShrAssign<usize> for BitSet {
131    #[inline]
132    fn shr_assign(&mut self, rhs: usize) {
133        let n = self.bits.len();
134        let k = rhs >> 6;
135        let d = rhs & 63;
136        if k >= n {
137            for x in self.bits.iter_mut() {
138                *x = 0;
139            }
140        } else {
141            if d == 0 {
142                for i in k..n {
143                    self.bits[i - k] = self.bits[i];
144                }
145            } else {
146                for i in k..n - 1 {
147                    self.bits[i - k] = (self.bits[i] >> d) | (self.bits[i + 1] << (64 - d));
148                }
149                self.bits[n - k - 1] = self.bits[n - 1] >> d;
150            }
151            for x in self.bits[n - k..].iter_mut() {
152                *x = 0;
153            }
154        }
155    }
156}
157impl Shr<usize> for BitSet {
158    type Output = Self;
159    #[inline]
160    fn shr(mut self, rhs: usize) -> Self::Output {
161        self >>= rhs;
162        self
163    }
164}
165impl<'a> BitOrAssign<&'a BitSet> for BitSet {
166    #[inline]
167    fn bitor_assign(&mut self, rhs: &'a Self) {
168        for (l, r) in self.bits.iter_mut().zip(rhs.bits.iter()) {
169            *l |= *r;
170        }
171        self.trim();
172    }
173}
174impl<'a> BitOr<&'a BitSet> for BitSet {
175    type Output = Self;
176    #[inline]
177    fn bitor(mut self, rhs: &'a Self) -> Self::Output {
178        self |= rhs;
179        self
180    }
181}
182impl<'b> BitOr<&'b BitSet> for &BitSet {
183    type Output = BitSet;
184    #[inline]
185    fn bitor(self, rhs: &'b BitSet) -> Self::Output {
186        let mut res = self.clone();
187        res |= rhs;
188        res
189    }
190}
191impl<'a> BitAndAssign<&'a BitSet> for BitSet {
192    #[inline]
193    fn bitand_assign(&mut self, rhs: &'a Self) {
194        for (l, r) in self.bits.iter_mut().zip(rhs.bits.iter()) {
195            *l &= *r;
196        }
197    }
198}
199impl<'a> BitAnd<&'a BitSet> for BitSet {
200    type Output = Self;
201    #[inline]
202    fn bitand(mut self, rhs: &'a Self) -> Self::Output {
203        self &= rhs;
204        self
205    }
206}
207impl<'b> BitAnd<&'b BitSet> for &BitSet {
208    type Output = BitSet;
209    #[inline]
210    fn bitand(self, rhs: &'b BitSet) -> Self::Output {
211        let mut res = self.clone();
212        res &= rhs;
213        res
214    }
215}
216impl<'a> BitXorAssign<&'a BitSet> for BitSet {
217    #[inline]
218    fn bitxor_assign(&mut self, rhs: &'a Self) {
219        for (l, r) in self.bits.iter_mut().zip(rhs.bits.iter()) {
220            *l ^= *r;
221        }
222        self.trim();
223    }
224}
225impl<'a> BitXor<&'a BitSet> for BitSet {
226    type Output = Self;
227    #[inline]
228    fn bitxor(mut self, rhs: &'a Self) -> Self::Output {
229        self ^= rhs;
230        self
231    }
232}
233impl<'b> BitXor<&'b BitSet> for &BitSet {
234    type Output = BitSet;
235    #[inline]
236    fn bitxor(self, rhs: &'b BitSet) -> Self::Output {
237        let mut res = self.clone();
238        res ^= rhs;
239        res
240    }
241}
242impl Not for BitSet {
243    type Output = Self;
244    #[inline]
245    fn not(mut self) -> Self::Output {
246        for x in self.bits.iter_mut() {
247            *x = !*x;
248        }
249        self.trim();
250        self
251    }
252}
253impl Not for &BitSet {
254    type Output = BitSet;
255    #[inline]
256    fn not(self) -> Self::Output {
257        !self.clone()
258    }
259}
competitive/data_structure/bitset.rs

competitive/data_structure/
bitset.rs
