Skip to content

Commit

Permalink
feat(docs): document Fenwick tree (#92)
Browse files Browse the repository at this point in the history
  • Loading branch information
scriptandcompile authored Jun 17, 2024
1 parent 189813b commit 76d17db
Showing 1 changed file with 101 additions and 35 deletions.
136 changes: 101 additions & 35 deletions src/data_structures/fenwick_tree.rs
Original file line number Diff line number Diff line change
@@ -1,40 +1,135 @@
use std::ops::{Add, AddAssign};

/// Fenwick Tree / Binary Indexed Tree
/// Consider we have an array arr[0 . . . n-1]. We would like to
/// A Fenwick Tree (also known as a Binary Indexed Tree) is a data structure
/// that can efficiently update elements and calculate prefix sums in a table of numbers.
///
/// If 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.Fenwick tree
/// 2. Modify the value of a specified element of the array arr[i] = x
/// where 0 <= i <= n-1.
///
/// A simple solution is to run a loop from 0 to i-1 and calculate the sum of the elements.
/// To update a value, simply do arr[i] = x. The first operation takes O(n) time and the
/// second operation takes O(1) time. Another simple solution is to create an extra array and
/// store the sum of the first i elements at the i-th index in this new array. The sum of a given
/// range can be calculated in O(1) time, but the update operation takes O(n) time now. This works
/// well if the number of query operations is large and the number of update operations is small.
///
/// The Fenwick tree provides a way to represent an array of numbers and perform two operations
/// in O(log n) time.
///
/// The two operations are:
/// 1. Update: Given a number x and an index i, we need to add x to the i-th element.
/// 2. Query: Given an index i, we need to calculate the prefix sum of the first i elements.
///
/// The Fenwick tree is represented as an array of n elements. The tree structure allows us to
/// perform efficient queries and updates. The tree is constructed in such a way that the value
/// of each node is the sum of the values of the nodes in its subtree.
///
pub struct FenwickTree<T: Add + AddAssign + Copy + Default> {
data: Vec<T>,
}

impl<T: Add<Output = T> + AddAssign + Copy + Default> FenwickTree<T> {
/// construct a new FenwickTree with given length
/// Create a new FenwickTree with length `len`
///
/// # Arguments
///
/// * `len` - The length of the FenwickTree
///
/// # Example
///
/// ```rust
/// use rust_algorithms::data_structures::FenwickTree;
///
/// let mut ft = FenwickTree::with_len(10);
/// ft.add(0, 1);
/// ft.add(1, 2);
///
/// assert_eq!(ft.prefix_sum(0), 1);
/// assert_eq!(ft.prefix_sum(1), 3);
/// ```
pub fn with_len(len: usize) -> Self {
FenwickTree {
data: vec![T::default(); len + 1],
}
}

/// add `val` to `idx`
/// Add `val` to the `i`-th element
///
/// # Arguments
///
/// * `i` - The index of the element to add `val` to
/// * `val` - The value to add to the `i`-th element
///
/// # Example
///
/// ```rust
/// use rust_algorithms::data_structures::FenwickTree;
///
/// let mut ft = FenwickTree::with_len(10);
/// ft.add(0, 3);
/// ft.add(1, 2);
///
/// assert_eq!(ft.prefix_sum(0), 3);
/// assert_eq!(ft.prefix_sum(1), 5);
/// ```
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);
}
}

/// get the sum of [0, i]
/// Get the prefix sum of the first `i` elements
///
/// # Arguments
///
/// * `i` - The index of the last element to calculate the prefix sum of
///
/// # Example
///
/// ```rust
/// use rust_algorithms::data_structures::FenwickTree;
///
/// 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);
/// ```
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);
}

res
}
}
Expand All @@ -44,32 +139,3 @@ const fn lowbit(x: usize) -> usize {
let x = x as isize;
(x & (-x)) as 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);
}
}

0 comments on commit 76d17db

Please sign in to comment.