haar_lib/ds/
succinct_bitvec.rs

1//! 簡潔ビットベクトル
2use std::ops::Range;
3
4const CHUNK_SIZE: usize = 256;
5const BLOCK_SIZE: usize = 64;
6const BLOCK_NUM: usize = CHUNK_SIZE / BLOCK_SIZE;
7
8/// 簡潔ビットベクトル
9#[derive(Clone)]
10pub struct SuccinctBitVec {
11    size: usize,
12    data: Vec<u64>,
13    blocks: Vec<[u8; BLOCK_NUM]>,
14    chunks: Vec<u32>,
15}
16
17impl SuccinctBitVec {
18    /// `Vec<bool>`から[`SuccinctBitVec`]を構築する。
19    pub fn new(b: Vec<bool>) -> Self {
20        let size = b.len();
21        let chunk_num = size.div_ceil(CHUNK_SIZE);
22        let mut data: Vec<u64> = vec![0; chunk_num * BLOCK_NUM + 1];
23
24        for (i, x) in b.into_iter().enumerate() {
25            if x {
26                let block_index = i / BLOCK_SIZE;
27                let index = i % BLOCK_SIZE;
28                data[block_index] |= 1 << index;
29            }
30        }
31
32        let mut chunks: Vec<u32> = vec![0; chunk_num + 1];
33        let mut blocks: Vec<[u8; BLOCK_NUM]> = vec![[0; BLOCK_NUM]; chunk_num + 1];
34
35        for (i, block_i) in blocks.iter_mut().take(chunk_num).enumerate() {
36            for j in 0..BLOCK_NUM - 1 {
37                block_i[j + 1] = block_i[j] + data[i * BLOCK_NUM + j].count_ones() as u8;
38            }
39
40            chunks[i + 1] = chunks[i]
41                + block_i[BLOCK_NUM - 1] as u32
42                + data[(i + 1) * BLOCK_NUM - 1].count_ones();
43        }
44
45        Self {
46            size,
47            data,
48            blocks,
49            chunks,
50        }
51    }
52
53    /// ビットベクトルの長さを返す。
54    pub fn len(&self) -> usize {
55        self.size
56    }
57
58    /// ビットベクトルの長さが`0`なら`true`を返す。
59    pub fn is_empty(&self) -> bool {
60        self.size == 0
61    }
62
63    /// [0, index) に含まれる`b`の個数
64    pub fn rank(&self, index: usize, b: bool) -> usize {
65        assert!(index <= self.size);
66
67        if b {
68            let chunk_pos = index / CHUNK_SIZE;
69            let block_pos = (index % CHUNK_SIZE) / BLOCK_SIZE;
70
71            let mask =
72                self.data[chunk_pos * BLOCK_NUM + block_pos] & ((1 << (index % BLOCK_SIZE)) - 1);
73
74            self.chunks[chunk_pos] as usize
75                + self.blocks[chunk_pos][block_pos] as usize
76                + mask.count_ones() as usize
77        } else {
78            index - self.rank(index, !b)
79        }
80    }
81
82    /// [l, r) に含まれる`b`の個数
83    pub fn count(&self, Range { start: l, end: r }: Range<usize>, b: bool) -> usize {
84        assert!(l <= r);
85        self.rank(r, b) - self.rank(l, b)
86    }
87
88    /// `index`番目のビットを返す。
89    pub fn access(&self, index: usize) -> u64 {
90        (self.data[index / BLOCK_SIZE] >> (index % BLOCK_SIZE)) & 1
91    }
92
93    /// nth(0-indexed)番目の`b`の位置
94    pub fn select(&self, nth: usize, b: bool) -> Option<usize> {
95        let nth = nth + 1;
96
97        if self.rank(self.size, b) < nth {
98            None
99        } else {
100            let mut lb: isize = -1;
101            let mut ub: isize = self.size as isize;
102            while (ub - lb).abs() > 1 {
103                let mid = (lb + ub) / 2;
104
105                if self.rank(mid as usize, b) >= nth {
106                    ub = mid;
107                } else {
108                    lb = mid;
109                }
110            }
111
112            Some(lb as usize)
113        }
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use rand::Rng;
121
122    #[test]
123    fn test_rank() {
124        let mut rng = rand::thread_rng();
125        let n = 100;
126        let b = (0..n).map(|_| rng.gen::<bool>()).collect::<Vec<_>>();
127
128        let s = SuccinctBitVec::new(b.clone());
129
130        for i in 0..=n {
131            let t = (0..i).filter(|&i| b[i]).count();
132            assert_eq!(s.rank(i, true), t);
133
134            let t = (0..i).filter(|&i| !b[i]).count();
135            assert_eq!(s.rank(i, false), t);
136        }
137    }
138
139    #[test]
140    fn test_count() {
141        let mut rng = rand::thread_rng();
142        let n = 100;
143        let b = (0..n).map(|_| rng.gen::<bool>()).collect::<Vec<_>>();
144
145        let s = SuccinctBitVec::new(b.clone());
146
147        for l in 0..=n {
148            for r in l..=n {
149                let t = (l..r).filter(|&i| b[i]).count();
150                assert_eq!(s.count(l..r, true), t);
151
152                let t = (l..r).filter(|&i| !b[i]).count();
153                assert_eq!(s.count(l..r, false), t);
154            }
155        }
156    }
157
158    #[test]
159    fn test_select() {
160        let mut rng = rand::thread_rng();
161        let n = 30;
162        let b = (0..n).map(|_| rng.gen::<bool>()).collect::<Vec<_>>();
163
164        let s = SuccinctBitVec::new(b.clone());
165
166        for i in 1..=n {
167            let t = (0..n).filter(|&i| b[i]).nth(i);
168            assert_eq!(s.select(i, true), t);
169
170            let t = (0..n).filter(|&i| !b[i]).nth(i);
171            assert_eq!(s.select(i, false), t);
172        }
173    }
174}
haar_lib/ds/succinct_bitvec.rs

haar_lib/ds/
succinct_bitvec.rs
