bit_vector/
lib.rs

1// N: O(n)
2// Nl: O(n log(n))
3// Nll: O(n log(log(n)))
4// L: O(log(n))
5// Ll: O(log(log(n)))
6// C: O(1), constant
7//
8// e.g. FooNllC: <O(n log(log(n))), O(1)> implementation for Foo.
9// In this context, it means <space complexity (in bits), query-time complexity>.
10
11use std::ops::Range;
12
13pub struct RankIndexNlC(Vec<usize>);
14
15impl RankIndexNlC {
16    pub fn new(a: &[bool]) -> Self {
17        let n = a.len();
18        let mut res = vec![0; n];
19        let mut count = 0;
20        for i in 0..n {
21            if a[i] {
22                count += 1;
23            }
24            res[i] = count;
25        }
26        Self(res)
27    }
28    pub fn rank<const X: bool>(&self, i: usize) -> usize {
29        if X { self.0[i] } else { i + 1 - self.0[i] }
30    }
31}
32
33pub struct SelectIndexNlC(Vec<usize>);
34
35impl SelectIndexNlC {
36    pub fn new<const X: bool>(a: &[bool]) -> Self {
37        let n = a.len();
38        let mut res = vec![0; n];
39        let mut count = 0;
40        for i in 0..n {
41            if a[i] == X {
42                res[count] = i;
43                count += 1;
44            }
45        }
46        Self(res)
47    }
48    pub fn select(&self, i: usize) -> usize { self.0[i] }
49}
50
51pub fn select_word<const X: bool>(mut w: u64, mut i: u32) -> u32 {
52    if !X {
53        w = !w;
54    }
55    let mut res = 0;
56    for lg2 in (0..6).rev() {
57        let len = 1 << lg2;
58        let mask = !(!0 << len);
59        let count = (w & mask).count_ones();
60        if count <= i {
61            w >>= len;
62            i -= count;
63            res += len;
64        }
65    }
66    res
67}
68
69pub struct Rs01DictNlC {
70    rank_index: RankIndexNlC,
71    select1_index: SelectIndexNlC,
72    select0_index: SelectIndexNlC,
73}
74
75impl Rs01DictNlC {
76    pub fn new(a: &[bool]) -> Self {
77        Self {
78            rank_index: RankIndexNlC::new(a),
79            select1_index: SelectIndexNlC::new::<true>(a),
80            select0_index: SelectIndexNlC::new::<false>(a),
81        }
82    }
83    fn rank<const X: bool>(&self, i: usize) -> usize {
84        self.rank_index.rank::<X>(i)
85    }
86    fn select<const X: bool>(&self, i: usize) -> usize {
87        if X {
88            self.select1_index.select(i)
89        } else {
90            self.select0_index.select(i)
91        }
92    }
93
94    pub fn rank1(&self, i: usize) -> usize { self.rank::<true>(i) }
95    pub fn rank0(&self, i: usize) -> usize { self.rank::<false>(i) }
96    pub fn select1(&self, i: usize) -> usize { self.select::<true>(i) }
97    pub fn select0(&self, i: usize) -> usize { self.select::<false>(i) }
98}
99
100struct RankIndexNLl {
101    block: Vec<usize>,
102    buf: Vec<u64>,
103}
104
105const W: usize = u64::BITS as usize;
106
107impl RankIndexNLl {
108    pub fn new(a: &[bool]) -> Self {
109        let len = a.len();
110        let n = (len + W - 1) / W;
111        let mut buf = vec![0_u64; n + 1];
112        for i in 0..len {
113            if a[i] {
114                buf[i / W] |= 1 << (i % W);
115            }
116        }
117        let block: Vec<_> = buf
118            .iter()
119            .map(|x| x.count_ones() as usize)
120            .scan(0, |acc, x| Some(std::mem::replace(acc, *acc + x)))
121            .collect();
122        Self { block, buf }
123    }
124    pub fn rank<const X: bool>(&self, i: usize) -> usize {
125        let i = i + 1;
126        let large = i / W;
127        let small = i % W;
128        let mini = self.buf[large] & !(!0 << small);
129        let count1 = self.block[large] + mini.count_ones() as usize;
130        if X { count1 } else { i - count1 }
131    }
132    pub fn rank_bisect<const X: bool>(
133        &self,
134        i: usize,
135        range: Range<usize>,
136    ) -> usize {
137        let mut lo = range.start / W;
138        let mut hi = (range.end + W - 1) / W;
139        while hi - lo > 1 {
140            let mid = lo + (hi - lo) / 2;
141            let count1 = self.block[mid];
142            let count = if X { count1 } else { mid * W - count1 };
143            *(if count <= i { &mut lo } else { &mut hi }) = mid;
144        }
145        let count1 = self.block[lo];
146        let count = if X { count1 } else { lo * W - count1 };
147        lo * W + select_word::<X>(self.buf[lo], (i - count) as _) as usize
148    }
149}
150
151struct SelectIndexNLl<const POPCNT: usize, const SPARSE_LEN: usize> {
152    inner: Vec<SelectIndexNLlInner>,
153}
154
155impl<const POPCNT: usize, const SPARSE_LEN: usize>
156    SelectIndexNLl<POPCNT, SPARSE_LEN>
157{
158    pub fn new<const X: bool>(a: &[bool]) -> Self {
159        let n = a.len();
160        let mut cur = vec![];
161        let mut res = vec![];
162        let mut start = 0;
163        for i in 0..n {
164            if a[i] == X {
165                cur.push(i);
166            }
167            if cur.len() >= POPCNT || i == n - 1 {
168                let tmp = std::mem::take(&mut cur);
169                if tmp.len() >= SPARSE_LEN {
170                    res.push(SelectIndexNLlInner::Sparse(tmp));
171                } else {
172                    res.push(SelectIndexNLlInner::Dense(start..i + 1));
173                }
174                start = i + 1;
175            }
176        }
177        Self { inner: res }
178    }
179
180    pub fn select<const X: bool>(&self, i: usize, r: &RankIndexNLl) -> usize {
181        self.inner[i / POPCNT].select::<X>(i, i % POPCNT, r)
182    }
183}
184
185enum SelectIndexNLlInner {
186    Sparse(Vec<usize>),
187    Dense(Range<usize>),
188}
189
190impl SelectIndexNLlInner {
191    fn select<const X: bool>(
192        &self,
193        i: usize,
194        i_rem: usize,
195        r: &RankIndexNLl,
196    ) -> usize {
197        match self {
198            Self::Sparse(pos) => pos[i_rem],
199            Self::Dense(Range { start, end }) => {
200                let lo = *start;
201                let hi = *end;
202                if r.rank::<X>(lo) > i {
203                    return lo;
204                }
205                r.rank_bisect::<X>(i, lo..hi)
206            }
207        }
208    }
209}
210
211const POPCNT: usize = 64; // log(n)
212const SPARSE_LEN: usize = 4096; // log(n)^2
213
214pub type Rs01DictNLl = Rs01DictNLlParam<POPCNT, SPARSE_LEN>;
215
216pub struct Rs01DictNLlParam<const POPCNT: usize, const SPARSE_LEN: usize> {
217    rank_index: RankIndexNLl,
218    select1_index: SelectIndexNLl<POPCNT, SPARSE_LEN>,
219    select0_index: SelectIndexNLl<POPCNT, SPARSE_LEN>,
220}
221
222impl<const POPCNT: usize, const SPARSE_LEN: usize>
223    Rs01DictNLlParam<POPCNT, SPARSE_LEN>
224{
225    pub fn new(a: &[bool]) -> Self {
226        Self {
227            rank_index: RankIndexNLl::new(a),
228            select1_index: SelectIndexNLl::new::<true>(a),
229            select0_index: SelectIndexNLl::new::<false>(a),
230        }
231    }
232    fn rank<const X: bool>(&self, i: usize) -> usize {
233        self.rank_index.rank::<X>(i)
234    }
235    fn select<const X: bool>(&self, i: usize) -> usize {
236        if X {
237            self.select1_index.select::<X>(i, &self.rank_index)
238        } else {
239            self.select0_index.select::<false>(i, &self.rank_index)
240        }
241    }
242
243    pub fn rank1(&self, i: usize) -> usize { self.rank::<true>(i) }
244    pub fn rank0(&self, i: usize) -> usize { self.rank::<false>(i) }
245    pub fn select1(&self, i: usize) -> usize { self.select::<true>(i) }
246    pub fn select0(&self, i: usize) -> usize { self.select::<false>(i) }
247}
248
249#[test]
250fn sanity_check() {
251    // % echo 0b_$(shuf -e {0,1}{0,1}{0,1}{0,1} | paste -sd _ -)_u64
252
253    let w = 0b_1100_1011_0010_1101_1001_1000_0001_0100_1111_1110_0000_0011_0111_0101_1010_0110_u64;
254    let mut wi = w;
255    for i in 0..32 {
256        let expected = wi.trailing_zeros();
257        let actual = select_word::<true>(w, i);
258        assert_eq!(actual, expected);
259        wi ^= wi & wi.wrapping_neg();
260    }
261}
262
263#[cfg(test)]
264macro_rules! bitvec {
265    ($lit:literal) => {
266        $lit.iter()
267            .filter(|&&b| matches!(b, b'0' | b'1'))
268            .map(|&b| b != b'0')
269            .collect::<Vec<_>>()
270    };
271}
272
273#[test]
274fn sanity_check_rank() {
275    let a = bitvec!(b"000 010 110 000; 111 001 000 011; 000 000 010 010");
276    let rs = Rs01DictNLlParam::<100, 100>::new(&a);
277    let expected1 = [
278        0, 0, 0, 0, 1, 1, 2, 3, 3, 3, 3, 3, 4, 5, 6, 6, 6, 7, 7, 7, 7, 7, 8, 9,
279        9, 9, 9, 9, 9, 9, 9, 10, 10, 10, 11, 11,
280    ];
281    let actual1: Vec<_> = (0..a.len()).map(|i| rs.rank1(i)).collect();
282    assert_eq!(actual1, expected1);
283
284    let expected0 = [
285        1, 2, 3, 4, 4, 5, 5, 5, 6, 7, 8, 9, 9, 9, 9, 10, 11, 11, 12, 13, 14,
286        15, 15, 15, 16, 17, 18, 19, 20, 21, 22, 22, 23, 24, 24, 25,
287    ];
288    let actual0: Vec<_> = (0..a.len()).map(|i| rs.rank0(i)).collect();
289    assert_eq!(actual0, expected0);
290}
291
292#[test]
293fn sanity_check_select_dense() {
294    let a = bitvec!(b"000 010 110; 000 111 001; 000 011 000");
295    let ones = a.iter().filter(|&&x| x).count();
296    let zeros = a.len() - ones;
297    let rs = Rs01DictNLlParam::<100, 100>::new(&a);
298    let expected1: Vec<_> = (0..a.len()).filter(|&i| a[i]).collect();
299    let expected0: Vec<_> = (0..a.len()).filter(|&i| !a[i]).collect();
300    let actual1: Vec<_> = (0..ones).map(|i| rs.select1(i)).collect();
301    let actual0: Vec<_> = (0..zeros).map(|i| rs.select0(i)).collect();
302    assert_eq!(actual1, expected1);
303    assert_eq!(actual0, expected0);
304}
305
306#[test]
307fn sanity_check_select_sparse() {
308    let a = bitvec!(b"001 010 000; 000 000 110");
309    let ones = a.iter().filter(|&&x| x).count();
310    let zeros = a.len() - ones;
311    let rs = Rs01DictNLlParam::<2, 0>::new(&a);
312    let expected1: Vec<_> = (0..a.len()).filter(|&i| a[i]).collect();
313    let expected0: Vec<_> = (0..a.len()).filter(|&i| !a[i]).collect();
314    let actual1: Vec<_> = (0..ones).map(|i| rs.select1(i)).collect();
315    let actual0: Vec<_> = (0..zeros).map(|i| rs.select0(i)).collect();
316    assert_eq!(actual1, expected1);
317    assert_eq!(actual0, expected0);
318}
bit_vector/lib.rs

bit_vector/
lib.rs
