ref: update Fenwick Tree implementation

- Add basic functinalities: `range_query`, `point_query`, `update`
and `set`
- Add docstring

src/data_structures/fenwick_tree.rs

@@ -1,46 +1,189 @@
-use std::ops::{Add, AddAssign};
+use std::ops::{Add, AddAssign, Sub, SubAssign};
 
-/// Fenwick Tree / Binary Indexed Tree
-///
-/// Consider we have an array `arr[0...n-1]`. We would like to:
-/// 1. Compute the sum of the first i elements.
-/// 2. Modify the value of a specified element of the array `arr[i] = x`, where `0 <= i <= n-1`.
-pub struct FenwickTree<T: Add + AddAssign + Copy + Default> {
+/// A Fenwick Tree (also known as a Binary Indexed Tree) that supports efficient
+/// prefix sum and range sum queries, as well as point updates.
+pub struct FenwickTree<T>
+where
+    T: Add<Output = T> + AddAssign + Sub<Output = T> + SubAssign + Copy + Default,
+{
+    /// Internal storage of the Fenwick Tree. The first element (index 0) is unused
+    /// to simplify index calculations, so the effective tree size is `data.len() - 1`.
     data: Vec<T>,
 }
 
-impl<T: Add<Output = T> + AddAssign + Copy + Default> FenwickTree<T> {
-    /// construct a new FenwickTree with given length
-    pub fn with_len(len: usize) -> Self {
+/// Enum representing the possible errors that can occur during FenwickTree operations.
+#[derive(Debug, PartialEq, Eq)]
+pub enum FenwickTreeError {
+    /// Error indicating that an index was out of the valid range.
+    IndexOutOfBounds,
+    /// Error indicating that a provided range was invalid (e.g., left > right).
+    InvalidRange,
+}
+
+impl<T> FenwickTree<T>
+where
+    T: Add<Output = T> + AddAssign + Sub<Output = T> + SubAssign + Copy + Default,
+{
+    /// Creates a new Fenwick Tree with a specified capacity.
+    ///
+    /// The tree will have `capacity + 1` elements, all initialized to the default
+    /// value of type `T`. The additional element allows for 1-based indexing internally.
+    ///
+    /// # Arguments
+    ///
+    /// * `capacity` - The number of elements the tree can hold (excluding the extra element).
+    ///
+    /// # Returns
+    ///
+    /// A new `FenwickTree` instance.
+    pub fn with_capacity(capacity: usize) -> Self {
         FenwickTree {
-            data: vec![T::default(); len + 1],
+            data: vec![T::default(); capacity + 1],
+        }
+    }
+
+    /// Updates the tree by adding a value to the element at a specified index.
+    ///
+    /// This operation also propagates the update to subsequent elements in the tree.
+    ///
+    /// # Arguments
+    ///
+    /// * `index` - The zero-based index where the value should be added.
+    /// * `value` - The value to add to the element at the specified index.
+    ///
+    /// # Returns
+    ///
+    /// A `Result` indicating success (`Ok`) or an error (`FenwickTreeError::IndexOutOfBounds`)
+    /// if the index is out of bounds.
+    pub fn update(&mut self, index: usize, value: T) -> Result<(), FenwickTreeError> {
+        if index >= self.data.len() - 1 {
+            return Err(FenwickTreeError::IndexOutOfBounds);
+        }
+
+        let mut idx = index + 1;
+        while idx < self.data.len() {
+            self.data[idx] += value;
+            idx += lowbit(idx);
         }
+
+        Ok(())
     }
 
-    /// add `val` to `idx`
-    pub fn add(&mut self, i: usize, val: T) {
-        assert!(i < self.data.len());
-        let mut i = i + 1;
-        while i < self.data.len() {
-            self.data[i] += val;
-            i += lowbit(i);
+    /// Computes the sum of elements from the start of the tree up to a specified index.
+    ///
+    /// This operation efficiently calculates the prefix sum using the tree structure.
+    ///
+    /// # Arguments
+    ///
+    /// * `index` - The zero-based index up to which the sum should be computed.
+    ///
+    /// # Returns
+    ///
+    /// A `Result` containing the prefix sum (`Ok(sum)`) or an error (`FenwickTreeError::IndexOutOfBounds`)
+    /// if the index is out of bounds.
+    pub fn prefix_query(&self, index: usize) -> Result<T, FenwickTreeError> {
+        if index >= self.data.len() - 1 {
+            return Err(FenwickTreeError::IndexOutOfBounds);
+        }
+
+        let mut idx = index + 1;
+        let mut result = T::default();
+        while idx > 0 {
+            result += self.data[idx];
+            idx -= lowbit(idx);
         }
+
+        Ok(result)
     }
 
-    /// get the sum of [0, i]
-    pub fn prefix_sum(&self, i: usize) -> T {
-        assert!(i < self.data.len());
-        let mut i = i + 1;
-        let mut res = T::default();
-        while i > 0 {
-            res += self.data[i];
-            i -= lowbit(i);
+    /// Computes the sum of elements within a specified range `[left, right]`.
+    ///
+    /// This operation calculates the range sum by performing two prefix sum queries.
+    ///
+    /// # Arguments
+    ///
+    /// * `left` - The zero-based starting index of the range.
+    /// * `right` - The zero-based ending index of the range.
+    ///
+    /// # Returns
+    ///
+    /// A `Result` containing the range sum (`Ok(sum)`) or an error:
+    /// * `FenwickTreeError::IndexOutOfBounds` if the right index is out of bounds.
+    /// * `FenwickTreeError::InvalidRange` if the left index is greater than the right index.
+    pub fn range_query(&self, left: usize, right: usize) -> Result<T, FenwickTreeError> {
+        if left > right || right >= self.data.len() - 1 {
+            return Err(FenwickTreeError::InvalidRange);
         }
-        res
+
+        let right_query = self.prefix_query(right)?;
+        let left_query = if left == 0 {
+            T::default()
+        } else {
+            self.prefix_query(left - 1)?
+        };
+
+        Ok(right_query - left_query)
+    }
+
+    /// Retrieves the value at a specific index by isolating it from the prefix sum.
+    ///
+    /// This operation determines the value at `index` by subtracting the prefix sum up to `index - 1`
+    /// from the prefix sum up to `index`.
+    ///
+    /// # Arguments
+    ///
+    /// * `index` - The zero-based index of the element to retrieve.
+    ///
+    /// # Returns
+    ///
+    /// A `Result` containing the value at the specified index (`Ok(value)`) or an error (`FenwickTreeError::IndexOutOfBounds`)
+    /// if the index is out of bounds.
+    pub fn point_query(&self, index: usize) -> Result<T, FenwickTreeError> {
+        if index >= self.data.len() - 1 {
+            return Err(FenwickTreeError::IndexOutOfBounds);
+        }
+
+        let index_query = self.prefix_query(index)?;
+        let prev_query = if index == 0 {
+            T::default()
+        } else {
+            self.prefix_query(index - 1)?
+        };
+
+        Ok(index_query - prev_query)
+    }
+
+    /// Sets the value at a specific index in the tree, updating the structure accordingly.
+    ///
+    /// This operation updates the value at `index` by computing the difference between the
+    /// desired value and the current value, then applying that difference using `update`.
+    ///
+    /// # Arguments
+    ///
+    /// * `index` - The zero-based index of the element to set.
+    /// * `value` - The new value to set at the specified index.
+    ///
+    /// # Returns
+    ///
+    /// A `Result` indicating success (`Ok`) or an error (`FenwickTreeError::IndexOutOfBounds`)
+    /// if the index is out of bounds.
+    pub fn set(&mut self, index: usize, value: T) -> Result<(), FenwickTreeError> {
+        self.update(index, value - self.point_query(index)?)
     }
 }
 
-/// get the lowest bit of `i`
+/// Computes the lowest set bit (rightmost `1` bit) of a number.
+///
+/// This function isolates the lowest set bit in the binary representation of `x`.
+/// It's used to navigate the Fenwick Tree by determining the next index to update or query.
+///
+/// # Arguments
+///
+/// * `x` - The input number whose lowest set bit is to be determined.
+///
+/// # Returns
+///
+/// The value of the lowest set bit in `x`.
 const fn lowbit(x: usize) -> usize {
     let x = x as isize;
     (x & (-x)) as usize
@@ -49,28 +192,67 @@ const fn lowbit(x: usize) -> usize {
 #[cfg(test)]
 mod tests {
     use super::*;
+
     #[test]
-    fn it_works() {
-        let mut ft = FenwickTree::with_len(10);
-        ft.add(0, 1);
-        ft.add(1, 2);
-        ft.add(2, 3);
-        ft.add(3, 4);
-        ft.add(4, 5);
-        ft.add(5, 6);
-        ft.add(6, 7);
-        ft.add(7, 8);
-        ft.add(8, 9);
-        ft.add(9, 10);
-        assert_eq!(ft.prefix_sum(0), 1);
-        assert_eq!(ft.prefix_sum(1), 3);
-        assert_eq!(ft.prefix_sum(2), 6);
-        assert_eq!(ft.prefix_sum(3), 10);
-        assert_eq!(ft.prefix_sum(4), 15);
-        assert_eq!(ft.prefix_sum(5), 21);
-        assert_eq!(ft.prefix_sum(6), 28);
-        assert_eq!(ft.prefix_sum(7), 36);
-        assert_eq!(ft.prefix_sum(8), 45);
-        assert_eq!(ft.prefix_sum(9), 55);
+    fn test_fenwick_tree() {
+        let mut fenwick_tree = FenwickTree::with_capacity(10);
+
+        assert_eq!(fenwick_tree.update(0, 5), Ok(()));
+        assert_eq!(fenwick_tree.update(1, 3), Ok(()));
+        assert_eq!(fenwick_tree.update(2, -2), Ok(()));
+        assert_eq!(fenwick_tree.update(3, 6), Ok(()));
+        assert_eq!(fenwick_tree.update(4, -4), Ok(()));
+        assert_eq!(fenwick_tree.update(5, 7), Ok(()));
+        assert_eq!(fenwick_tree.update(6, -1), Ok(()));
+        assert_eq!(fenwick_tree.update(7, 2), Ok(()));
+        assert_eq!(fenwick_tree.update(8, -3), Ok(()));
+        assert_eq!(fenwick_tree.update(9, 4), Ok(()));
+        assert_eq!(fenwick_tree.set(3, 10), Ok(()));
+        assert_eq!(fenwick_tree.point_query(3), Ok(10));
+        assert_eq!(fenwick_tree.set(5, 0), Ok(()));
+        assert_eq!(fenwick_tree.point_query(5), Ok(0));
+        assert_eq!(
+            fenwick_tree.update(10, 11),
+            Err(FenwickTreeError::IndexOutOfBounds)
+        );
+        assert_eq!(
+            fenwick_tree.set(10, 11),
+            Err(FenwickTreeError::IndexOutOfBounds)
+        );
+
+        assert_eq!(fenwick_tree.prefix_query(0), Ok(5));
+        assert_eq!(fenwick_tree.prefix_query(1), Ok(8));
+        assert_eq!(fenwick_tree.prefix_query(2), Ok(6));
+        assert_eq!(fenwick_tree.prefix_query(3), Ok(16));
+        assert_eq!(fenwick_tree.prefix_query(4), Ok(12));
+        assert_eq!(fenwick_tree.prefix_query(5), Ok(12));
+        assert_eq!(fenwick_tree.prefix_query(6), Ok(11));
+        assert_eq!(fenwick_tree.prefix_query(7), Ok(13));
+        assert_eq!(fenwick_tree.prefix_query(8), Ok(10));
+        assert_eq!(fenwick_tree.prefix_query(9), Ok(14));
+        assert_eq!(
+            fenwick_tree.prefix_query(10),
+            Err(FenwickTreeError::IndexOutOfBounds)
+        );
+
+        assert_eq!(fenwick_tree.range_query(0, 4), Ok(12));
+        assert_eq!(fenwick_tree.range_query(3, 7), Ok(7));
+        assert_eq!(fenwick_tree.range_query(2, 5), Ok(4));
+        assert_eq!(
+            fenwick_tree.range_query(4, 3),
+            Err(FenwickTreeError::InvalidRange)
+        );
+        assert_eq!(
+            fenwick_tree.range_query(2, 10),
+            Err(FenwickTreeError::InvalidRange)
+        );
+
+        assert_eq!(fenwick_tree.point_query(0), Ok(5));
+        assert_eq!(fenwick_tree.point_query(4), Ok(-4));
+        assert_eq!(fenwick_tree.point_query(9), Ok(4));
+        assert_eq!(
+            fenwick_tree.point_query(10),
+            Err(FenwickTreeError::IndexOutOfBounds)
+        );
     }
 }