1#![forbid(unsafe_code)]
72#![feature(test)]
73
74use std::ops::{Deref, DerefMut};
75
76mod fixed_size_tree;
77mod growing_tree;
78
79pub use fixed_size_tree::FixedSizeFenwickTree;
80pub use growing_tree::GrowingFenwickTree;
81
82pub mod prelude {
84 pub use crate::FenwickTreeValue;
85 pub use crate::fixed_size_tree::FixedSizeFenwickTree;
86 pub use crate::growing_tree::GrowingFenwickTree;
87 pub use crate::FenwickTree;
88 pub use crate::TreeError;
89}
90
91fn least_significant_bit(idx: usize) -> usize {
92 let int_idx = idx as i32;
93 (int_idx & -int_idx) as usize
94}
95
96pub trait FenwickTreeValue:
98 Default + Clone + core::cmp::PartialEq
100{
101 fn store_value(&mut self, other: &Self);
102 fn substract(self, other: Self) -> Self;
103}
104
105impl<T> FenwickTreeValue for T
106where T: Default + Copy + std::ops::AddAssign
108 + std::ops::Sub<Output = Self>
109 + core::cmp::PartialEq
110{
111 fn store_value(&mut self, other: &Self) {
112 *self += *other
113 }
114
115 fn substract(self, other: Self) -> Self {
116 self - other
117 }
118}
119
120pub trait FenwickTree {
122 type Value: FenwickTreeValue;
123
124 fn query(&self, idx: usize) -> Result<Self::Value, TreeError>;
132
133 fn update(&mut self, idx: usize, value: Self::Value) -> Result<(), TreeError>;
141
142 fn range_query(&self, from: usize, to: usize) -> Result<Self::Value, TreeError> {
151 let from_sum = self.query(from)?;
152 let to_sum = self.query(to)?;
153 Ok(to_sum.substract(from_sum))
154 }
155}
156
157#[derive(Debug, Clone, Copy)]
161enum TreeIndex {
162 Internal { val: usize },
163 External { val: usize },
164}
165
166#[derive(Debug, PartialEq)]
167pub enum TreeError {
168 IndexOutOfBounds( usize )
169}
170
171impl TreeIndex {
172
173 fn to_internal(self) -> Self {
174 match self {
175 TreeIndex::Internal { val: _ } => self,
176 TreeIndex::External { val } => TreeIndex::Internal { val: val + 1 },
177 }
178 }
179
180 fn to_external(self) -> Result<Self, String> {
181 match self {
182 TreeIndex::Internal { val } => {
183 if val == 0 {
184 return Err("Index is out of bounds.".to_string());
185 }
186 Ok(TreeIndex::External { val: val - 1 })
187 }
188 TreeIndex::External { val: _ } => Ok(self),
189 }
190 }
191
192 fn lsb_descending(self) -> LeastSignificantBitDescentingChain {
195 LeastSignificantBitDescentingChain {
196 idx: self.to_internal(),
197 }
198 }
199
200 fn lsb_ascending(self, upper_bound: usize) -> LeastSignificantBitAscendingChain {
203 LeastSignificantBitAscendingChain {
204 idx: self.to_internal(),
205 max: upper_bound,
206 }
207 }
208
209 fn is_power_of_2(self) -> bool {
210 let idx = *self;
211 idx.is_power_of_two()
212 }
213
214}
215
216impl From<usize> for TreeIndex {
217 fn from(value: usize) -> Self {
218 Self::External { val: value }
219 }
220}
221
222impl Deref for TreeIndex {
223 type Target = usize;
224
225 fn deref(&self) -> &Self::Target {
226 match self {
227 TreeIndex::External { val } => val,
228 TreeIndex::Internal { val } => val,
229 }
230 }
231}
232
233impl PartialEq for TreeIndex {
234 fn eq(&self, other: &Self) -> bool {
235 match (self, other) {
236 (Self::Internal { val: l_val }, Self::Internal { val: r_val }) => l_val == r_val,
237 (Self::External { val: l_val }, Self::External { val: r_val }) => l_val == r_val,
238 _ => false,
239 }
240 }
241}
242
243impl DerefMut for TreeIndex {
244 fn deref_mut(&mut self) -> &mut Self::Target {
245 match self {
246 TreeIndex::External { val } => val,
247 TreeIndex::Internal { val } => val,
248 }
249 }
250}
251
252struct LeastSignificantBitDescentingChain {
255 idx: TreeIndex,
256}
257
258impl Iterator for LeastSignificantBitDescentingChain {
259 type Item = TreeIndex;
260
261 fn next(&mut self) -> Option<Self::Item> {
262 if *self.idx == 0 {
263 return None;
264 }
265 let res = TreeIndex::Internal { val: *self.idx };
267 *self.idx -= least_significant_bit(*self.idx);
268 Some(res)
269 }
270}
271
272struct LeastSignificantBitAscendingChain {
275 idx: TreeIndex,
276 max: usize,
277}
278
279impl Iterator for LeastSignificantBitAscendingChain {
280 type Item = TreeIndex;
281
282 fn next(&mut self) -> Option<Self::Item> {
283 if *self.idx > self.max {
284 return None;
285 }
286 let res = TreeIndex::Internal { val: *self.idx };
288 *self.idx += least_significant_bit(*self.idx);
289 Some(res)
290 }
291}
292
293#[cfg(test)]
294mod tests {
295
296 use pretty_assertions::assert_eq;
297
298 use crate::{least_significant_bit, TreeIndex};
299
300 fn to_internal_index_vec(indexes: &[usize]) -> Vec<TreeIndex> {
301 indexes
302 .into_iter()
303 .map(|i| TreeIndex::Internal { val: *i })
304 .collect::<Vec<TreeIndex>>()
305 }
306
307 #[test]
308 fn test_index_transform_from_internal_to_external_with_error() {
309 let idx = TreeIndex::Internal { val: 0 };
310 idx.to_external().expect_err("Index is out of bounds.");
311 }
312
313 #[test]
314 fn test_index_transform_from_internal_to_external() {
315 for val in 1..100 {
316 let idx = TreeIndex::Internal { val: val };
317 assert_eq!(
318 idx.to_external().unwrap(),
319 TreeIndex::External { val: val - 1 }
320 );
321 }
322 }
323
324 #[test]
325 fn test_index_transform_from_external_to_internal() {
326 for val in 0..100 {
327 let idx = TreeIndex::External { val: val };
328 assert_eq!(idx.to_internal(), TreeIndex::Internal { val: val + 1 });
329 }
330 }
331
332 #[test]
333 fn test_index_transform_to_itseld() {
334 for val in 0..100 {
335 let idx = TreeIndex::External { val: val };
336 assert_eq!(idx.to_external().unwrap(), TreeIndex::External { val });
337 }
338
339 for val in 0..100 {
340 let idx = TreeIndex::Internal { val: val };
341 assert_eq!(idx.to_internal(), TreeIndex::Internal { val: val });
342 }
343 }
344
345 #[test]
346 fn test_ascending_lsb_chain() {
347 let idx: TreeIndex = 0.into();
348 assert_eq!(
349 idx.lsb_ascending(64).collect::<Vec<TreeIndex>>(),
350 to_internal_index_vec(&[1, 2, 4, 8, 16, 32, 64])
351 );
352
353 let idx: TreeIndex = 1.into();
354 assert_eq!(
355 idx.lsb_ascending(64).collect::<Vec<TreeIndex>>(),
356 to_internal_index_vec(&[2, 4, 8, 16, 32, 64])
357 );
358
359 let idx: TreeIndex = 6.into();
360 assert_eq!(
361 idx.lsb_ascending(64).collect::<Vec<TreeIndex>>(),
362 to_internal_index_vec(&[7, 8, 16, 32, 64])
363 );
364
365 let idx: TreeIndex = 6.into();
366 assert_eq!(idx.lsb_ascending(0).collect::<Vec<TreeIndex>>(), vec![]);
367 }
368
369 #[test]
370 fn test_descending_lsb_chain() {
371 let idx: TreeIndex = 5.into();
372 assert_eq!(idx, TreeIndex::External { val: 5 });
373 assert_eq!(
374 idx.lsb_descending().collect::<Vec<TreeIndex>>(),
375 to_internal_index_vec(&[6, 4])
376 );
377
378 let idx: TreeIndex = 4.into();
379 assert_eq!(
380 idx.lsb_descending().collect::<Vec<TreeIndex>>(),
381 to_internal_index_vec(&[5, 4])
382 );
383
384 let idx = TreeIndex::Internal { val: 3 };
385 assert_eq!(
386 idx.lsb_descending().collect::<Vec<TreeIndex>>(),
387 to_internal_index_vec(&[3, 2])
388 );
389
390 let idx = TreeIndex::Internal { val: 12 };
391 assert_eq!(
392 idx.lsb_descending().collect::<Vec<TreeIndex>>(),
393 to_internal_index_vec(&[12, 8])
394 );
395 }
396
397 #[test]
398 fn test_lsb() {
399 assert_eq!(least_significant_bit(12), 4)
400 }
401
402 #[test]
403 fn test_bitwise_op() {
404 assert_eq!(12usize.next_power_of_two(), 16);
405 assert_eq!(12usize.next_power_of_two() >> 1, 8);
406 }
407}