bitree/lib.rs
1//! A binary indexed tree (Fenwick tree) for efficient prefix sums.
2//!
3//! A [`BITree<T>`] maintains prefix sums over a mutable sequence of values.
4//! Both point updates and prefix queries run in *O*(log *n*) time, and the
5//! whole structure lives in a single [`Vec<T>`] the same length as the
6//! logical array.
7//!
8//! This crate is `no_std`-compatible; it depends only on [`alloc`]. Values are
9//! generic over any `T` implementing `AddAssign<&T>` — and, where a method
10//! requires subtraction, `SubAssign<&T>`. `T: Copy` is not required, so
11//! arbitrary-precision integers and similar types are supported.
12//!
13//! # Examples
14//!
15//! ```
16//! use bitree::BITree;
17//!
18//! let mut bitree = BITree::from_iter([1, 6, 3, 9, 2]);
19//!
20//! // Prefix sums in O(log n); `prefix_sum(i)` covers `[0..i)`.
21//! assert_eq!(bitree.prefix_sum(3), 10);
22//! assert_eq!(bitree.prefix_sum(5), 21);
23//!
24//! // Point update.
25//! bitree.add_at(1, 5);
26//! assert_eq!(bitree.prefix_sum(3), 15);
27//!
28//! // Grow and shrink like a `Vec`.
29//! bitree.push(4);
30//! bitree.pop();
31//!
32//! // Recover the original values.
33//! assert_eq!(Vec::from(bitree), vec![1, 11, 3, 9, 2]);
34//! ```
35//!
36//! # Crate features
37//!
38//! - **`serde`** — derive [`Serialize`] and [`Deserialize`] for [`BITree<T>`].
39//!
40//! [`Serialize`]: https://docs.rs/serde/latest/serde/trait.Serialize.html
41//! [`Deserialize`]: https://docs.rs/serde/latest/serde/trait.Deserialize.html
42
43#![no_std]
44
45extern crate alloc;
46
47use alloc::vec::Vec;
48use core::ops::{AddAssign, SubAssign};
49
50#[cfg(feature = "serde")]
51use serde::{Deserialize, Serialize};
52
53/// A binary indexed tree (Fenwick tree) over a sequence of `T` values.
54///
55/// Conceptually, a `BITree<T>` represents an array `a[0..n]` while internally
56/// storing partial sums that allow both prefix-sum queries and point updates
57/// in *O*(log *n*). The internal buffer has the same length as the logical
58/// array, so no memory is wasted relative to a plain [`Vec<T>`].
59///
60/// See the [crate-level documentation](crate) for an overview and examples.
61#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
62#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Hash)]
63pub struct BITree<T> {
64 inner: Vec<T>,
65}
66
67impl<T> BITree<T> {
68 /// Constructs a new, empty `BITree<T>`.
69 ///
70 /// The tree will not allocate until values are pushed into it.
71 ///
72 /// # Examples
73 ///
74 /// ```
75 /// use bitree::BITree;
76 ///
77 /// let bitree: BITree<i32> = BITree::new();
78 /// assert!(bitree.is_empty());
79 /// ```
80 #[inline]
81 pub const fn new() -> Self {
82 Self { inner: Vec::new() }
83 }
84
85 /// Constructs a new `BITree<T>` of length `n`, filled with `T::default()`.
86 ///
87 /// The initial capacity is exactly `n`.
88 ///
89 /// # Examples
90 ///
91 /// ```
92 /// use bitree::BITree;
93 ///
94 /// let bitree = BITree::<usize>::new_zeros(5);
95 /// assert_eq!(bitree.len(), 5);
96 /// assert_eq!(bitree.prefix_sum(5), 0);
97 /// ```
98 #[inline]
99 pub fn new_zeros(n: usize) -> Self
100 where
101 T: Default,
102 {
103 let mut inner = Vec::with_capacity(n);
104 for _ in 0..n {
105 inner.push(T::default());
106 }
107
108 Self { inner }
109 }
110
111 /// Constructs a new, empty `BITree<T>` with at least the specified capacity.
112 ///
113 /// The tree will be able to hold at least `capacity` elements without
114 /// reallocating. This method is allowed to allocate for more elements
115 /// than `capacity`. If `capacity` is zero, the tree will not allocate.
116 ///
117 /// # Examples
118 ///
119 /// ```
120 /// use bitree::BITree;
121 ///
122 /// let mut bitree: BITree<i32> = BITree::with_capacity(10);
123 /// assert!(bitree.is_empty());
124 ///
125 /// for i in 0..10 {
126 /// bitree.push(i);
127 /// }
128 /// assert_eq!(bitree.len(), 10);
129 /// ```
130 #[inline]
131 pub fn with_capacity(capacity: usize) -> Self {
132 Self {
133 inner: Vec::with_capacity(capacity),
134 }
135 }
136
137 /// Returns `true` if the tree contains no elements.
138 ///
139 /// # Examples
140 ///
141 /// ```
142 /// use bitree::BITree;
143 ///
144 /// let mut bitree = BITree::new();
145 /// assert!(bitree.is_empty());
146 ///
147 /// bitree.push(1);
148 /// assert!(!bitree.is_empty());
149 /// ```
150 #[inline]
151 pub const fn is_empty(&self) -> bool {
152 self.inner.is_empty()
153 }
154
155 /// Returns the number of elements in the tree.
156 ///
157 /// # Examples
158 ///
159 /// ```
160 /// use bitree::BITree;
161 ///
162 /// let bitree = BITree::from_iter([1, 2, 3]);
163 /// assert_eq!(bitree.len(), 3);
164 /// ```
165 #[inline]
166 pub const fn len(&self) -> usize {
167 self.inner.len()
168 }
169
170 /// Removes the last element from the tree, returning `true` if one was
171 /// removed and `false` if the tree was already empty.
172 ///
173 /// # Examples
174 ///
175 /// ```
176 /// use bitree::BITree;
177 ///
178 /// let mut bitree = BITree::from_iter([1, 6, 3, 9]);
179 ///
180 /// assert_eq!(bitree.pop(), true);
181 /// assert_eq!(bitree.prefix_sum(3), 10); // sum of remaining [1, 6, 3]
182 ///
183 /// assert_eq!(bitree.pop(), true);
184 /// assert_eq!(bitree.prefix_sum(2), 7); // sum of remaining [1, 6]
185 ///
186 /// bitree.pop();
187 /// bitree.pop();
188 /// assert_eq!(bitree.pop(), false); // already empty
189 /// ```
190 #[inline(always)]
191 pub fn pop(&mut self) -> bool {
192 self.inner.pop().is_some()
193 }
194
195 #[inline(always)]
196 fn walk_prefix<F: FnMut(&mut T, &T)>(&self, index: usize, sum: &mut T, mut op: F) {
197 assert!(index < self.inner.len() + 1);
198
199 let mut current_idx = index;
200 while current_idx > 0 {
201 op(sum, &self.inner[current_idx - 1]);
202 current_idx &= current_idx - 1;
203 }
204 }
205
206 #[inline(always)]
207 fn walk_update<F: FnMut(&mut T, &T)>(&mut self, index: usize, diff: T, mut op: F) {
208 assert!(index < self.len());
209
210 let mut current_idx = index;
211 while let Some(value) = self.inner.get_mut(current_idx) {
212 op(value, &diff);
213 current_idx |= current_idx + 1;
214 }
215 }
216}
217
218impl<T> Default for BITree<T> {
219 /// Creates an empty `BITree<T>`.
220 #[inline]
221 fn default() -> Self {
222 Self::new()
223 }
224}
225
226impl<T: for<'a> AddAssign<&'a T>> From<Vec<T>> for BITree<T> {
227 /// Builds a `BITree<T>` from a [`Vec`] of values, reusing its allocation.
228 ///
229 /// Runs in *O*(*n*).
230 ///
231 /// # Examples
232 ///
233 /// ```
234 /// use bitree::BITree;
235 ///
236 /// let values = vec![1, 6, 3, 9, 2];
237 /// let bitree = BITree::from(values);
238 /// assert_eq!(bitree.prefix_sum(4), 19);
239 /// ```
240 #[inline]
241 fn from(mut inner: Vec<T>) -> Self {
242 let n = inner.len();
243 rebuild(&mut inner, 0..n, |p, c| *p += c);
244 BITree { inner }
245 }
246}
247
248impl<T: for<'a> SubAssign<&'a T>> From<BITree<T>> for Vec<T> {
249 /// Recovers the original values as a [`Vec`], reusing the tree's allocation.
250 ///
251 /// Runs in *O*(*n*).
252 ///
253 /// # Examples
254 ///
255 /// ```
256 /// use bitree::BITree;
257 ///
258 /// let bitree = BITree::from(vec![1, 6, 3, 9, 2]);
259 /// assert_eq!(Vec::from(bitree), vec![1, 6, 3, 9, 2]);
260 /// ```
261 #[inline]
262 fn from(mut bitree: BITree<T>) -> Self {
263 let n = bitree.inner.len();
264 rebuild(&mut bitree.inner, (0..n).rev(), |p, c| *p -= c);
265 bitree.inner
266 }
267}
268
269impl<T: for<'a> AddAssign<&'a T>> FromIterator<T> for BITree<T> {
270 /// Builds a `BITree<T>` from the values yielded by an iterator.
271 ///
272 /// Runs in *O*(*n*).
273 ///
274 /// # Examples
275 ///
276 /// ```
277 /// use bitree::BITree;
278 ///
279 /// let bitree = BITree::from_iter([1, 6, 3, 9, 2]);
280 /// assert_eq!(bitree.prefix_sum(5), 21);
281 /// ```
282 #[inline]
283 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
284 Self::from(iter.into_iter().collect::<Vec<_>>())
285 }
286}
287
288impl<T: for<'a> SubAssign<&'a T>> IntoIterator for BITree<T> {
289 type Item = T;
290 type IntoIter = alloc::vec::IntoIter<T>;
291
292 /// Consumes the tree and yields the original values in order.
293 ///
294 /// The returned iterator implements both [`DoubleEndedIterator`] and
295 /// [`ExactSizeIterator`], so forward and reverse traversal both run in
296 /// *O*(1) per element after an *O*(*n*) setup that inverts the Fenwick
297 /// build.
298 ///
299 /// # Examples
300 ///
301 /// ```
302 /// use bitree::BITree;
303 ///
304 /// let bitree = BITree::from_iter([1, 6, 3, 9, 2]);
305 /// let collected: Vec<_> = bitree.into_iter().collect();
306 /// assert_eq!(collected, vec![1, 6, 3, 9, 2]);
307 /// ```
308 ///
309 /// Reverse iteration works too:
310 ///
311 /// ```
312 /// use bitree::BITree;
313 ///
314 /// let bitree = BITree::from_iter([1, 6, 3, 9, 2]);
315 /// let reversed: Vec<_> = bitree.into_iter().rev().collect();
316 /// assert_eq!(reversed, vec![2, 9, 3, 6, 1]);
317 /// ```
318 #[inline]
319 fn into_iter(self) -> Self::IntoIter {
320 Vec::from(self).into_iter()
321 }
322}
323
324#[inline(always)]
325fn rebuild<T, I, F>(inner: &mut [T], indices: I, mut op: F)
326where
327 I: IntoIterator<Item = usize>,
328 F: FnMut(&mut T, &T),
329{
330 let n = inner.len();
331 let ptr = inner.as_mut_ptr();
332 for i in indices {
333 let parent = i | (i + 1);
334 if parent < n {
335 // SAFETY:
336 // - i < parent < n, so both offsets are in-bounds of `inner`.
337 // - parent != i, so the &mut and & never alias.
338 // - `ptr` is derived from a valid &mut [T] that outlives the loop.
339 unsafe {
340 let child = &*ptr.add(i);
341 let parent_ref = &mut *ptr.add(parent);
342 op(parent_ref, child);
343 }
344 }
345 }
346}
347
348impl<T: for<'a> AddAssign<&'a T> + for<'a> SubAssign<&'a T>> BITree<T> {
349 /// Adds the prefix sum of `[0..index)` into `sum`.
350 ///
351 /// When `index` is `0`, `sum` is left unchanged.
352 ///
353 /// # Panics
354 ///
355 /// Panics if `index > self.len()`.
356 ///
357 /// # Examples
358 ///
359 /// ```
360 /// use bitree::BITree;
361 ///
362 /// let bitree = BITree::from_iter([1, 6, 3, 9, 2]);
363 ///
364 /// let mut running = 100;
365 /// bitree.add_prefix_sum(3, &mut running);
366 /// assert_eq!(running, 110);
367 /// ```
368 #[inline]
369 pub fn add_prefix_sum(&self, index: usize, sum: &mut T) {
370 self.walk_prefix(index, sum, |s, v| *s += v);
371 }
372
373 /// Subtracts the prefix sum of `[0..index)` from `sum`.
374 ///
375 /// When `index` is `0`, `sum` is left unchanged.
376 ///
377 /// # Panics
378 ///
379 /// Panics if `index > self.len()`.
380 ///
381 /// # Examples
382 ///
383 /// ```
384 /// use bitree::BITree;
385 ///
386 /// let bitree = BITree::from_iter([1, 6, 3, 9, 2]);
387 ///
388 /// let mut running = 100;
389 /// bitree.sub_prefix_sum(3, &mut running);
390 /// assert_eq!(running, 90);
391 /// ```
392 #[inline]
393 pub fn sub_prefix_sum(&self, index: usize, sum: &mut T) {
394 self.walk_prefix(index, sum, |s, v| *s -= v);
395 }
396
397 /// Returns the prefix sum of `[0..index)`.
398 ///
399 /// Equivalent to starting from `T::default()` and calling
400 /// [`add_prefix_sum`](Self::add_prefix_sum). The prefix sum at `index = 0`
401 /// is `T::default()`.
402 ///
403 /// # Panics
404 ///
405 /// Panics if `index > self.len()`.
406 ///
407 /// # Examples
408 ///
409 /// ```
410 /// use bitree::BITree;
411 ///
412 /// let bitree = BITree::from_iter([1, 6, 3, 9, 2]);
413 ///
414 /// assert_eq!(bitree.prefix_sum(0), 0);
415 /// assert_eq!(bitree.prefix_sum(1), 1);
416 /// assert_eq!(bitree.prefix_sum(3), 10);
417 /// assert_eq!(bitree.prefix_sum(5), 21);
418 /// ```
419 #[inline]
420 pub fn prefix_sum(&self, index: usize) -> T
421 where
422 T: Default,
423 {
424 let mut sum = T::default();
425 self.add_prefix_sum(index, &mut sum);
426 sum
427 }
428
429 /// Adds `diff` to the value at `index`.
430 ///
431 /// # Panics
432 ///
433 /// Panics if `index >= self.len()`.
434 ///
435 /// # Examples
436 ///
437 /// ```
438 /// use bitree::BITree;
439 ///
440 /// let mut bitree = BITree::from_iter([1, 6, 3, 9, 2]);
441 /// // logical array: [1, 6, 3, 9, 2]
442 ///
443 /// bitree.add_at(1, 5);
444 /// // logical array: [1, 11, 3, 9, 2]
445 ///
446 /// assert_eq!(bitree.prefix_sum(3), 15);
447 /// ```
448 #[inline]
449 pub fn add_at(&mut self, index: usize, diff: T) {
450 self.walk_update(index, diff, |v, d| *v += d);
451 }
452
453 /// Subtracts `diff` from the value at `index`.
454 ///
455 /// # Panics
456 ///
457 /// Panics if `index >= self.len()`.
458 ///
459 /// # Examples
460 ///
461 /// ```
462 /// use bitree::BITree;
463 ///
464 /// let mut bitree = BITree::from_iter([1, 6, 3, 9, 2]);
465 /// // logical array: [1, 6, 3, 9, 2]
466 ///
467 /// bitree.sub_at(3, 4);
468 /// // logical array: [1, 6, 3, 5, 2]
469 ///
470 /// assert_eq!(bitree.prefix_sum(5), 17);
471 /// ```
472 #[inline]
473 pub fn sub_at(&mut self, index: usize, diff: T) {
474 self.walk_update(index, diff, |v, d| *v -= d);
475 }
476
477 /// Appends a value to the end of the tree.
478 ///
479 /// Runs in amortised *O*(log *n*).
480 ///
481 /// # Examples
482 ///
483 /// ```
484 /// use bitree::BITree;
485 ///
486 /// let mut bitree = BITree::from_iter([1, 6, 3]);
487 /// bitree.push(9);
488 /// bitree.push(2);
489 ///
490 /// assert_eq!(bitree.prefix_sum(4), 19);
491 /// assert_eq!(bitree.prefix_sum(5), 21);
492 /// ```
493 #[inline]
494 pub fn push(&mut self, mut value: T) {
495 let n = self.inner.len();
496 for i in 0..n.trailing_ones() {
497 let child = n & !(1 << i);
498 value += &self.inner[child];
499 }
500 self.inner.push(value);
501 }
502}
503
504impl<T: for<'a> AddAssign<&'a T> + for<'a> SubAssign<&'a T> + PartialOrd> BITree<T> {
505 /// Walks the tree to find the largest `pos` such that
506 /// `prefix_sum(pos) <= target`, subtracting the consumed segment sums
507 /// from `*remainder` along the way.
508 ///
509 /// The caller supplies the initial `target` in `*remainder`. After the
510 /// call, `*remainder` holds `target - prefix_sum(pos)`. A remainder of
511 /// `T::default()` means `target` landed exactly on the boundary
512 /// `prefix_sum(pos)`; any positive remainder falls strictly inside the
513 /// slot that starts at `pos`.
514 ///
515 /// This is the lower-level primitive behind
516 /// [`binary_search`](Self::binary_search); use that method when you only
517 /// need the slot index and don't care about the remainder.
518 ///
519 /// Equality is decided via the trichotomy of [`PartialOrd`], so
520 /// [`PartialEq`] is not required. Values that are incomparable with
521 /// `target` (for example `f64::NAN`) are the caller's responsibility.
522 ///
523 /// # Examples
524 ///
525 /// ```
526 /// use bitree::BITree;
527 ///
528 /// let bitree = BITree::from_iter([1, 6, 3, 9, 2]);
529 ///
530 /// // 9 lies between prefix_sum(2) = 7 and prefix_sum(3) = 10.
531 /// let mut remaining = 9;
532 /// let idx = bitree.sub_binary_search(&mut remaining);
533 /// assert_eq!((idx, remaining), (2, 2));
534 ///
535 /// // Exact boundary: 7 == prefix_sum(2), so the walk advances past it.
536 /// let mut remaining = 7;
537 /// let idx = bitree.sub_binary_search(&mut remaining);
538 /// assert_eq!((idx, remaining), (2, 0));
539 /// ```
540 pub fn sub_binary_search(&self, remainder: &mut T) -> usize {
541 let n = self.inner.len();
542 let mut pos = 0;
543
544 let mut mask = n.checked_ilog2().map_or(0, |log| 1 << log);
545
546 while mask > 0 {
547 let next = pos + mask;
548 if next <= n {
549 let value = &self.inner[next - 1];
550 if !(*remainder < *value) {
551 pos = next;
552 *remainder -= value;
553 }
554 }
555 mask >>= 1;
556 }
557
558 pos
559 }
560
561 /// Binary-searches the conceptual prefix-sum slice
562 /// `[prefix_sum(0), prefix_sum(1), ..., prefix_sum(n)]` for `target`.
563 ///
564 /// The behaviour mirrors [`slice::binary_search`]:
565 ///
566 /// - `Ok(k)` if `prefix_sum(k) == target`.
567 /// - `Err(k)` if `target` would be inserted at position `k` to keep the
568 /// slice sorted — that is, `prefix_sum(k - 1) < target < prefix_sum(k)`,
569 /// with `Err(0)` when `target` is below `prefix_sum(0)` and
570 /// `Err(n + 1)` when `target` is above `prefix_sum(n)`.
571 ///
572 /// Equality is decided via the trichotomy of [`PartialOrd`], so
573 /// [`PartialEq`] is not required. Values that are incomparable with
574 /// `target` (for example `f64::NAN`) are the caller's responsibility.
575 ///
576 /// # Examples
577 ///
578 /// Looking up sums in an integer tree:
579 ///
580 /// ```
581 /// use bitree::BITree;
582 ///
583 /// let bitree = BITree::from_iter([1, 6, 3, 9, 2]);
584 /// // prefix sums: 0, 1, 7, 10, 19, 21
585 ///
586 /// assert_eq!(bitree.binary_search(0), Ok(0));
587 /// assert_eq!(bitree.binary_search(7), Ok(2));
588 /// assert_eq!(bitree.binary_search(21), Ok(5));
589 ///
590 /// assert_eq!(bitree.binary_search(6), Err(2));
591 /// assert_eq!(bitree.binary_search(9), Err(3));
592 /// assert_eq!(bitree.binary_search(22), Err(6));
593 /// ```
594 ///
595 /// An empty tree still has the single prefix sum `prefix_sum(0) = 0`:
596 ///
597 /// ```
598 /// use bitree::BITree;
599 ///
600 /// let empty = BITree::<f64>::new();
601 ///
602 /// assert!(empty.is_empty());
603 /// assert_eq!(empty.binary_search(-1.0), Err(0));
604 /// assert_eq!(empty.binary_search(0.0), Ok(0));
605 /// assert_eq!(empty.binary_search(1.0), Err(1));
606 /// ```
607 #[inline]
608 pub fn binary_search(&self, mut target: T) -> Result<usize, usize>
609 where
610 T: Default,
611 {
612 let pos = self.sub_binary_search(&mut target);
613 let zero = T::default();
614 if target < zero {
615 Err(pos)
616 } else if zero < target {
617 Err(pos + 1)
618 } else {
619 Ok(pos)
620 }
621 }
622}
623
624#[cfg(test)]
625mod tests {
626 extern crate std;
627 use super::BITree;
628 use alloc::vec;
629 use alloc::vec::Vec;
630
631 #[test]
632 fn test_new() {
633 let lengths: [usize; 5] = [1, 6, 3, 9, 2];
634 let expected_index = vec![1, 7, 3, 19, 2];
635 let actual_index = BITree::from_iter(lengths);
636 assert_eq!(expected_index, actual_index.inner);
637
638 let n = 5;
639 let tree = BITree::<usize>::new_zeros(5);
640 assert_eq!(tree.len(), n);
641 assert!(!tree.is_empty());
642 assert_eq!(tree.prefix_sum(0), 0);
643 assert_eq!(tree.prefix_sum(3), 0);
644 assert_eq!(tree.prefix_sum(5), 0);
645 }
646
647 #[test]
648 fn test_prefix_sum() {
649 let lengths = [1, 6, 3, 9, 2];
650 let bitree = BITree::from_iter(lengths);
651
652 let cases: Vec<(usize, usize)> = vec![(0, 0), (1, 1), (2, 7), (3, 10), (4, 19), (5, 21)];
653 // The prefix sum up until the zeroth element is 0, since there is nothing before it
654 // The prefix sum up until an index larger than the length is undefined, since every
655 // element after the length - 1 is undefined
656 cases
657 .into_iter()
658 .for_each(|(idx, expected_sum)| assert_eq!(bitree.prefix_sum(idx), expected_sum))
659 }
660
661 #[test]
662 fn test_update_index() {
663 let lengths = [1, 6, 3, 9, 2];
664 let mut bitree = BITree::from_iter(lengths);
665
666 let cases: Vec<(usize, usize)> = vec![(0, 2), (1, 8), (2, 3), (3, 20), (4, 2)];
667
668 bitree.add_at(0, 1);
669
670 cases
671 .into_iter()
672 .for_each(|(idx, expected_value)| assert_eq!(bitree.inner[idx], expected_value))
673 }
674
675 #[test]
676 fn test_binary_search() {
677 let lengths = [1, 6, 3, 9, 2];
678 let bitree = BITree::from_iter(lengths);
679 // prefix sums: 0, 1, 7, 10, 19, 21
680
681 let cases: Vec<(usize, Result<usize, usize>)> = vec![
682 (0, Ok(0)),
683 (1, Ok(1)),
684 (7, Ok(2)),
685 (10, Ok(3)),
686 (19, Ok(4)),
687 (21, Ok(5)),
688 (6, Err(2)),
689 (9, Err(3)),
690 (18, Err(4)),
691 (20, Err(5)),
692 (22, Err(6)),
693 ];
694
695 cases
696 .into_iter()
697 .for_each(|(target, expected)| assert_eq!(bitree.binary_search(target), expected))
698 }
699
700 #[test]
701 #[ntest::timeout(1000)]
702 fn test_zero_array() {
703 // regression: a tree containing only 0 used to loop endlessly
704 let f0: BITree<usize> = BITree::from_iter([0]);
705 assert_eq!(f0.prefix_sum(0), 0);
706 // prefix sums: [0, 0]; searching for 1 falls past the end.
707 assert_eq!(f0.binary_search(1), Err(2));
708 let mut remaining = 1usize;
709 assert_eq!(f0.sub_binary_search(&mut remaining), 1);
710 assert_eq!(remaining, 1);
711 }
712
713 #[test]
714 fn test_sub_binary_search_empty() {
715 let bitree: BITree<usize> = BITree::new();
716 let mut remaining = 5;
717 assert_eq!(bitree.sub_binary_search(&mut remaining), 0);
718 assert_eq!(remaining, 5);
719 assert_eq!(bitree.binary_search(0usize), Ok(0));
720 assert_eq!(bitree.binary_search(5usize), Err(1));
721 }
722
723 #[test]
724 fn test_sub_binary_search_single() {
725 let bitree = BITree::from_iter([7usize]);
726 // prefix sums: 0, 7
727 let cases: Vec<(usize, (usize, usize))> =
728 vec![(0, (0, 0)), (1, (0, 1)), (7, (1, 0)), (8, (1, 1))];
729 cases.into_iter().for_each(|(target, expected)| {
730 let mut remaining = target;
731 let idx = bitree.sub_binary_search(&mut remaining);
732 assert_eq!((idx, remaining), expected, "target={}", target);
733 });
734 }
735
736 #[test]
737 fn test_sub_binary_search_power_of_two_len() {
738 // length 4 exercises a tree whose root mask equals n
739 let bitree = BITree::from_iter([2usize, 3, 5, 7]);
740 // prefix sums: 0, 2, 5, 10, 17
741 let cases: Vec<(usize, (usize, usize))> = vec![
742 (0, (0, 0)),
743 (2, (1, 0)), // boundary
744 (3, (1, 1)),
745 (5, (2, 0)), // boundary
746 (6, (2, 1)),
747 (10, (3, 0)), // boundary
748 (11, (3, 1)),
749 (17, (4, 0)), // boundary: total sum
750 (18, (4, 1)), // exceeds total
751 ];
752 cases.into_iter().for_each(|(target, expected)| {
753 let mut remaining = target;
754 let idx = bitree.sub_binary_search(&mut remaining);
755 assert_eq!((idx, remaining), expected, "target={}", target);
756 });
757 }
758
759 #[test]
760 fn test_sub_binary_search_uniform_seven() {
761 // length 7 (odd, non-power-of-two) with uniform values makes expected
762 // results easy to reason about: prefix_sum(k) = k.
763 let bitree = BITree::from_iter([1usize; 7]);
764 let cases: Vec<(usize, (usize, usize))> = vec![
765 (0, (0, 0)),
766 (1, (1, 0)),
767 (2, (2, 0)),
768 (3, (3, 0)),
769 (4, (4, 0)),
770 (5, (5, 0)),
771 (6, (6, 0)),
772 (7, (7, 0)),
773 (8, (7, 1)), // exceeds total
774 ];
775 cases.into_iter().for_each(|(target, expected)| {
776 let mut remaining = target;
777 let idx = bitree.sub_binary_search(&mut remaining);
778 assert_eq!((idx, remaining), expected, "target={}", target);
779 });
780 }
781
782 #[test]
783 fn test_sub_binary_search_exceeds_total() {
784 let bitree = BITree::from_iter([1usize, 6, 3, 9, 2]);
785 // total sum = 21, n = 5
786 let mut remaining = 100;
787 assert_eq!(bitree.sub_binary_search(&mut remaining), 5);
788 assert_eq!(remaining, 100 - 21);
789 }
790
791 #[test]
792 fn test_push_empty() {
793 let mut bitree = BITree::new();
794 bitree.push(5);
795 assert_eq!(bitree.inner, vec![5]);
796 assert_eq!(bitree.prefix_sum(1), 5);
797 }
798
799 #[test]
800 fn test_push_sequence() {
801 let mut bitree = BITree::new();
802 let values = [1, 6, 3, 9, 2];
803 let expected_sums = vec![(1, 1), (2, 7), (3, 10), (4, 19), (5, 21)];
804
805 for &v in values.iter() {
806 bitree.push(v);
807 }
808
809 expected_sums
810 .into_iter()
811 .for_each(|(idx, expected_sum)| assert_eq!(bitree.prefix_sum(idx), expected_sum));
812 }
813
814 #[test]
815 fn test_push_after_initialization() {
816 let mut bitree = BITree::from_iter([1, 6, 3].into_iter());
817 bitree.push(9);
818 bitree.push(2);
819
820 let expected_sums = vec![(1, 1), (2, 7), (3, 10), (4, 19), (5, 21)];
821 expected_sums
822 .into_iter()
823 .for_each(|(idx, expected_sum)| assert_eq!(bitree.prefix_sum(idx), expected_sum));
824 }
825
826 #[test]
827 fn test_pop_empty() {
828 let mut bitree: BITree<usize> = BITree::new();
829 assert_eq!(bitree.pop(), false);
830 }
831
832 #[test]
833 fn test_pop_single() {
834 let mut bitree = BITree::from_iter([5].into_iter());
835 assert_eq!(bitree.pop(), true);
836 assert!(bitree.is_empty());
837 }
838
839 #[test]
840 fn test_pop_sequence() {
841 let mut bitree = BITree::from_iter([1, 6, 3, 9, 2].into_iter());
842 assert_eq!(bitree.pop(), true);
843 assert_eq!(bitree.pop(), true);
844 assert_eq!(bitree.pop(), true);
845
846 assert_eq!(bitree.prefix_sum(1), 1);
847 assert_eq!(bitree.prefix_sum(2), 7);
848 }
849
850 #[test]
851 fn test_push_pop_alternating() {
852 let mut bitree = BITree::new();
853
854 bitree.push(1);
855 bitree.push(6);
856 assert_eq!(bitree.pop(), true);
857 bitree.push(3);
858 assert_eq!(bitree.pop(), true);
859 bitree.push(9);
860 bitree.push(2);
861 assert_eq!(bitree.pop(), true);
862
863 assert_eq!(bitree.prefix_sum(1), 1);
864 assert_eq!(bitree.prefix_sum(2), 10);
865 }
866
867 #[test]
868 fn test_zero_handling() {
869 let mut bitree = BITree::new();
870 bitree.push(0);
871 bitree.push(0);
872 assert_eq!(bitree.pop(), true);
873 assert_eq!(bitree.prefix_sum(1), 0);
874 }
875
876 #[test]
877 fn test_negative_values() {
878 let mut bitree: BITree<i32> = BITree::new();
879 bitree.push(-1);
880 bitree.push(2);
881 bitree.push(-3);
882
883 assert_eq!(bitree.pop(), true);
884 assert_eq!(bitree.prefix_sum(2), 1);
885 }
886}