dsalgo/
bit_array.rs

1// TODO: currently, this is implemented as dynamic sized.
2// but bit array should have static size.
3const K: usize = 6; // chunk size = 2^K
4const M: usize = (1 << K) - 1;
5
6fn bucket(index: usize) -> usize {
7    index >> K
8}
9
10fn point(index: usize) -> usize {
11    index & M
12}
13
14fn value(index: usize) -> usize {
15    1 << point(index)
16}
17
18#[derive(Clone)]
19
20pub struct BitArray(Vec<usize>);
21
22impl BitArray {
23    pub fn new(size: usize) -> Self {
24        BitArray(vec![0; (size + M) >> K])
25    }
26
27    pub fn set(
28        &mut self,
29        i: usize,
30    ) {
31        self.0[bucket(i)] |= value(i);
32    }
33
34    pub fn reset(
35        &mut self,
36        i: usize,
37    ) {
38        self.0[bucket(i)] &= !value(i);
39    }
40
41    pub fn flip(
42        &mut self,
43        i: usize,
44    ) {
45        self.0[bucket(i)] ^= value(i);
46    }
47
48    pub fn is_set(
49        &self,
50        i: usize,
51    ) -> bool {
52        self.0[bucket(i)] >> point(i) & 1 == 1
53    }
54}
55
56impl From<&[bool]> for BitArray {
57    fn from(is_set: &[bool]) -> Self {
58        let mut a = BitArray::new(is_set.len());
59
60        for (i, &ok) in is_set.iter().enumerate() {
61            if ok {
62                a.set(i)
63            }
64        }
65
66        a
67    }
68}
69
70use std::ops::*;
71
72fn min_len(
73    lhs: &BitArray,
74    rhs: &BitArray,
75) -> usize {
76    lhs.0.len().min(rhs.0.len())
77}
78
79fn max_len(
80    lhs: &BitArray,
81    rhs: &BitArray,
82) -> usize {
83    lhs.0.len().max(rhs.0.len())
84}
85
86impl BitAnd for BitArray {
87    type Output = Self;
88
89    fn bitand(
90        mut self,
91        rhs: Self,
92    ) -> Self::Output {
93        for i in 0..min_len(&self, &rhs) {
94            self.0[i] &= rhs.0[i];
95        }
96
97        self
98    }
99}
100
101impl BitArray {
102    pub fn popcount(&self) -> u64 {
103        self.0.iter().map(|x| x.count_ones() as u64).sum()
104    }
105}
106
107// TODO:
108// &, |,  <<, >>, ==
109#[cfg(test)]
110
111mod tests {
112
113    #[test]
114
115    fn test() {}
116}