fenwick_tree/
tree.rs

1use std::ops::{AddAssign, Bound, RangeBounds, SubAssign};
2
3use crate::errors::{AddError, SumError};
4
5/// An implementation of the binary indexed tree (Fenwick tree) data structure.
6///
7/// The tree is backed by a simple vec of the fixed size where each item is
8/// responsible for storing cumulative sum of some range, allowing to perform
9/// queries and updates in _O_(log _n_) time.
10pub struct FenwickTree<I>
11where
12    I: Default + Copy + AddAssign + SubAssign,
13{
14    tree: Vec<I>,
15}
16
17impl<I> FenwickTree<I>
18where
19    I: Default + Copy + AddAssign + SubAssign,
20{
21    /// Constructs a new Fenwick tree with the specified `len` with each element set as
22    /// `I::default()`.
23    ///
24    /// The vector is initialized with `vec![I::default(); len]`.
25    ///
26    /// # Panics
27    ///
28    /// Vector initialization may panic if `len` is too big.
29    pub fn with_len(len: usize) -> Self {
30        Self {
31            tree: vec![I::default(); len],
32        }
33    }
34
35    /// A length of the backing vector of the tree.
36    pub fn len(&self) -> usize {
37        self.tree.len()
38    }
39
40    /// A partial sum of the specified `bounds`.
41    ///
42    /// Complexity: _O_(log _n_).
43    ///
44    /// Note that `sum` for empty range (a range `(i..j)` where `i >= j`) is `0`.
45    ///
46    /// Also, `bounds` are converted into a pair `(start, end)` that represents `[start,
47    /// end)` range. This means that boundary case `tree.sum(i..=usize::MAX)` fallbacks to `tree.sum(0..usize::MAX)`.
48    /// However, in practice, it's not possible to construct such a big tree ([`Vec`] panics with
49    /// `capacity overflow`).
50    pub fn sum(&self, bounds: impl RangeBounds<usize>) -> Result<I, SumError> {
51        let len = self.len();
52
53        let mut sum = I::default();
54        let mut start = start(bounds.start_bound());
55        let mut end = end(bounds.end_bound(), len);
56
57        if start >= len || end > len {
58            return Err(SumError::OutOfRange {
59                bounds: as_pair(bounds),
60                len,
61            });
62        }
63
64        while end > start {
65            sum += self.tree[end - 1];
66            end = prev(end);
67        }
68
69        while start > end {
70            sum -= self.tree[start - 1];
71            start = prev(start);
72        }
73
74        return Ok(sum);
75
76        // As inclusive.
77        fn start(bound: Bound<&usize>) -> usize {
78            match bound {
79                Bound::Excluded(&usize::MAX) => usize::MAX,
80                Bound::Excluded(x) => *x + 1,
81                Bound::Included(x) => *x,
82                Bound::Unbounded => 0,
83            }
84        }
85
86        // As exclusive.
87        fn end(bound: Bound<&usize>, len: usize) -> usize {
88            match bound {
89                Bound::Included(&usize::MAX) => usize::MAX,
90                Bound::Included(x) => *x + 1,
91                Bound::Excluded(x) => *x,
92                Bound::Unbounded => len,
93            }
94        }
95
96        fn as_pair(bounds: impl RangeBounds<usize>) -> (Bound<usize>, Bound<usize>) {
97            let start = cloned(bounds.start_bound());
98            let end = cloned(bounds.end_bound());
99
100            return (start, end);
101
102            // From nightly Rust (https://doc.rust-lang.org/std/ops/enum.Bound.html#method.cloned).
103            fn cloned(x: Bound<&usize>) -> Bound<usize> {
104                match x {
105                    Bound::Unbounded => Bound::Unbounded,
106                    Bound::Included(x) => Bound::Included(*x),
107                    Bound::Excluded(x) => Bound::Excluded(*x),
108                }
109            }
110        }
111    }
112
113    /// Updates the value at `i` by `delta`.
114    ///
115    /// Complexity: _O_(log _n_).
116    pub fn add(&mut self, mut i: usize, delta: I) -> Result<(), AddError> {
117        let size = self.len();
118
119        if i >= size {
120            return Err(AddError::IndexOutOfRange { index: i, size });
121        }
122
123        while i < size {
124            self.tree[i] += delta;
125            i = next(i);
126        }
127
128        Ok(())
129    }
130}
131
132/// Flips first trailing `1` in the binary representation of the `i`. Same as `i - (i & (-i))` (see
133/// crate docs).
134///
135/// This allows fast calculating of prefix sums:
136///     - call `i = prev(i)` until `i` is greater than 0
137///     - access sums by `i - 1`
138///
139/// This function assumes that indexing is one-based, hence we access sums by `i - 1`.
140/// However, it's worth to note that zero-based solution (`i & (i + 1)`) produces less cleaner code
141/// because to iterate we need to call `i = prev(i) - 1`, which involves additional checks when `i`
142/// is of `usize` (decrement may result in panic).
143const fn prev(i: usize) -> usize {
144    i & (i - 1)
145}
146
147/// Flips first trailing `0` in the binary representation of the `i`.
148///
149/// In the same way as with `prev`, `i = next(i)` allows traversing the array but in the opposite
150/// direction.
151/// However, unlike `prev`, this function assumes that indexing is zero-based, hence we access sums by `i`.
152const fn next(i: usize) -> usize {
153    i | (i + 1)
154}
fenwick_tree/tree.rs

fenwick_tree/
tree.rs
