fenwick_bit_tree/
growing_tree.rs

1use crate::{FenwickTree, FenwickTreeValue, TreeError, TreeIndex};
2
3pub struct GrowingFenwickTree<T> {
4    data: Vec<T>,
5}
6
7impl<T: FenwickTreeValue> GrowingFenwickTree<T> {
8    pub fn new(size: usize) -> Self {
9        Self {
10            data: vec![T::default(); size + 1],
11        }
12    }
13
14    fn size(&self) -> usize {
15        self.data.len()
16    }
17
18    fn resize(&mut self, idx: &TreeIndex) -> Result<(), TreeError> {
19        let size_before_resize = self.size();
20
21        // TODO: resize should grow to the closest including power of 2
22        self.data.resize(*idx.to_internal() + 1, T::default());
23
24        if size_before_resize <= 1 {
25            return Ok(());
26        }
27
28        let highest_index_before_resize = TreeIndex::Internal {
29            val: size_before_resize - 1,
30        };
31
32        let first_new_index = TreeIndex::Internal {
33            val: size_before_resize,
34        };
35
36        let aggregate_from = if highest_index_before_resize.is_power_of_2() {
37            TreeIndex::Internal { val: 0 }
38        } else {
39            first_new_index
40                .lsb_descending()
41                .skip(1)
42                .chain([TreeIndex::Internal { val: 0 }])
43                .next()
44                .unwrap()
45        };
46
47        let sum_from = aggregate_from
48            .to_external()
49            .map_or(Ok(T::default()), |idx| self.query(*idx))?;
50
51        let sum_till = highest_index_before_resize
52            .to_external()
53            .map_or(Ok(T::default()), |idx| self.query(*idx))?;
54
55        let value = sum_till.substract(sum_from);
56
57        for data_position in highest_index_before_resize
58            .lsb_ascending(self.size() - 1)
59            .skip(1)
60        {
61            let data_position = data_position.to_internal();
62            self[data_position].store_value(&value);
63        }
64
65        Ok(())
66    }
67}
68
69impl<T> std::ops::Index<TreeIndex> for GrowingFenwickTree<T> {
70    type Output = T;
71
72    fn index(&self, index: TreeIndex) -> &Self::Output {
73        &self.data[*index.to_internal()]
74    }
75}
76
77impl<T> std::ops::IndexMut<TreeIndex> for GrowingFenwickTree<T> {
78    fn index_mut(&mut self, index: TreeIndex) -> &mut Self::Output {
79        &mut self.data[*index.to_internal()]
80    }
81}
82
83impl<T: FenwickTreeValue> FenwickTree for GrowingFenwickTree<T> {
84    type Value = T;
85
86    fn query(&self, idx: usize) -> Result<T, TreeError> {
87        let mut idx: TreeIndex = idx.into();
88
89        if self.size() <= *idx.to_internal() {
90            idx = TreeIndex::Internal {
91                val: self.size() - 1,
92            }
93        }
94
95        let mut res = Self::Value::default();
96
97        for data_position in idx.lsb_descending() {
98            let data_position = data_position.to_internal();
99            res.store_value(&self[data_position]);
100        }
101
102        Ok(res)
103    }
104
105    fn update(&mut self, idx: usize, value: Self::Value) -> Result<(), TreeError> {
106        let idx: TreeIndex = idx.into();
107
108        if *idx.to_internal() > self.size() - 1 {
109            self.resize(&idx)?
110        }
111
112        for data_position in idx.lsb_ascending(self.size() - 1) {
113            let data_position = data_position.to_internal();
114            self[data_position].store_value(&value);
115        }
116
117        Ok(())
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use rand::seq::SliceRandom;
124    use rand::Rng;
125
126    use crate::growing_tree::GrowingFenwickTree;
127    use crate::FenwickTree;
128
129    #[test]
130    fn empty_tree_query() {
131        let tree = GrowingFenwickTree::<i32>::new(0);
132        assert!(tree.query(0).is_ok_and(|val| val == 0));
133        assert!(tree.query(1).is_ok_and(|val| val == 0));
134    }
135
136    #[test]
137    fn one_element_tree_query() {
138        let tree = GrowingFenwickTree::<i32>::new(1);
139        assert!(tree.query(0).is_ok_and(|val| val == 0));
140        assert!(tree.query(1).is_ok_and(|val| val == 0));
141    }
142
143    #[test]
144    fn test_no_upper_bound_error_is_raised() {
145        let tree = GrowingFenwickTree::<i32>::new(0);
146        assert_eq!(tree.query(100).unwrap(), 0);
147        assert_eq!(tree.range_query(10, 100).unwrap(), 0);
148    }
149
150    #[test]
151    fn tree_grows_one_by_one() {
152        let mut tree = GrowingFenwickTree::<i32>::new(1);
153        tree.update(3, 1).unwrap();
154        assert_eq!(tree.query(3).unwrap(), 1);
155
156        tree.update(0, 1).unwrap();
157        assert_eq!(tree.query(3).unwrap(), 2);
158    }
159
160    #[test]
161    fn tree_suddenly_grows_much_bigger() {
162        let mut tree = GrowingFenwickTree::<i32>::new(2);
163        tree.update(0, 1).unwrap();
164        assert_eq!(tree.query(0).unwrap(), 1);
165
166        tree.update(1, 1).unwrap();
167        assert_eq!(tree.query(1).unwrap(), 2);
168
169        tree.update(7, 0).unwrap();
170        assert_eq!(tree.query(7).unwrap(), 2);
171    }
172
173    #[test]
174    fn simple_tree_generation_with_queries() {
175        let mut tree = GrowingFenwickTree::<i32>::new(11);
176        for i in 0..32 {
177            if let Err(_) = tree.update(i, 1) {
178                assert!(false)
179            }
180        }
181        assert_eq!(tree.query(3).unwrap(), 4); // points at [0, 1, 2, 3, 4]
182        assert_eq!(tree.query(0).unwrap(), 1);
183        assert_eq!(tree.query(31).unwrap(), 32);
184    }
185
186    #[test]
187    fn test_range_queries() {
188        let mut tree = GrowingFenwickTree::<i32>::new(0);
189        for i in 0..=29 {
190            if let Err(_) = tree.update(i, 1) {
191                assert!(false)
192            }
193        }
194
195        match tree.range_query(10, 20) {
196            Ok(10) => assert!(true),
197            _ => assert!(false),
198        }
199        match tree.range_query(8, 29) {
200            Ok(21) => assert!(true),
201            _ => assert!(false),
202        }
203    }
204
205    #[test]
206    fn update_existent_value() {
207        let mut tree = GrowingFenwickTree::<i32>::new(0);
208        for _i in 0..32 {
209            if let Err(_) = tree.update(0, 1) {
210                assert!(false)
211            }
212        }
213        let res = tree.query(0).unwrap();
214        assert_eq!(res, 32);
215    }
216
217    #[test]
218    fn random_100_point_data() {
219        let size = 100;
220        let mut input = vec![];
221        let mut rng = rand::thread_rng();
222
223        for _i in 0..size {
224            input.push((rng.gen::<f32>() * 100.0) as i32);
225        }
226
227        let mut tree = GrowingFenwickTree::<i32>::new(0);
228        for i in 0..size {
229            if let Err(_) = tree.update(i, *input.get(i).unwrap()) {
230                assert!(false)
231            }
232        }
233
234        let mut sum = 0;
235        for i in 0..size {
236            sum += *input.get(i).unwrap();
237
238            if let Ok(res) = tree.query(i) {
239                assert_eq!(res, sum);
240            } else {
241                assert!(false)
242            }
243        }
244    }
245
246    #[test]
247    fn random_100_point_data_with_random_update_order() {
248        let size = 100;
249        let mut input = vec![];
250        let mut rng = rand::thread_rng();
251
252        for _i in 0..size {
253            input.push((rng.gen::<f32>() * 100.0) as i32);
254        }
255
256        let mut tree = GrowingFenwickTree::<i32>::new(size);
257
258        let mut random_indexes: Vec<usize> = (0..size).collect();
259        random_indexes.shuffle(&mut rng);
260        for i in random_indexes {
261            if let Err(_) = tree.update(i, *input.get(i).unwrap()) {
262                assert!(false)
263            }
264        }
265
266        let mut sum = 0;
267        for i in 0..size {
268            sum += *input.get(i).unwrap();
269            if let Ok(res) = tree.query(i) {
270                assert_eq!(res, sum);
271            } else {
272                assert!(false);
273            }
274        }
275    }
276
277    #[test]
278    fn random_100_point_data_with_random_update_order_with_intermediate_asserts() {
279        let size = 100;
280        let mut input = vec![];
281        let mut rng = rand::thread_rng();
282
283        for _i in 0..size {
284            input.push((rng.gen::<f32>() * 100.0) as i32);
285        }
286
287        let mut tree = GrowingFenwickTree::<i32>::new(size);
288
289        let mut random_indexes: Vec<usize> = (0..size).collect();
290        random_indexes.shuffle(&mut rng);
291        for i in random_indexes {
292            let sum_before_update = tree.query(i).unwrap();
293            let value_to_update = *input.get(i).unwrap();
294            if let Err(_) = tree.update(i, value_to_update) {
295                assert!(false)
296            }
297            let sum_after_update = tree.query(i).unwrap();
298            assert_eq!(sum_after_update - sum_before_update, value_to_update)
299        }
300
301        let mut sum = 0;
302        for i in 0..size {
303            sum += *input.get(i).unwrap();
304
305            if let Ok(res) = tree.query(i) {
306                assert_eq!(res, sum);
307            } else {
308                assert!(false)
309            }
310        }
311    }
312}
fenwick_bit_tree/growing_tree.rs

fenwick_bit_tree/
growing_tree.rs
