Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added documentation for fenwick_tree.rs, also cleaned up the descript… #92

Merged
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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);
}
}
Loading