diff --git a/src/data_structures/fenwick_tree.rs b/src/data_structures/fenwick_tree.rs index 5066d2c..8ee308d 100644 --- a/src/data_structures/fenwick_tree.rs +++ b/src/data_structures/fenwick_tree.rs @@ -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 { data: Vec, } impl + AddAssign + Copy + Default> FenwickTree { - /// 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 } } @@ -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); - } -}