fenwick/
lib.rs

1//! Fenwick Tree (Binary Indexed Tree)
2//!
3//! Fenwick Tree is a data structure that can efficiently update elements
4//! and calculate prefix sums in a table of numbers.
5//!
6//! Given an array of `n` numbers, the prefix sum up to index `i` is the
7//! sum of the first `i` elements. Fenwick Tree supports both update and
8//! prefix-sum queries in `O(log n)` time by storing partial sums in a tree
9//! whose node ranges are determined by the least significant bit of each
10//! index.
11
12use std::ops::{AddAssign, Sub};
13
14/// Fenwick Tree (Binary Indexed Tree).
15///
16/// `T` is the numeric type stored in the tree. It must be `Copy`,
17/// `Default` (the additive identity, typically `0`), and support
18/// `+=` and `-`.
19#[derive(Debug, Clone)]
20pub struct FenwickTree<T> {
21    size: usize,
22    tree: Vec<T>,
23}
24
25impl<T> FenwickTree<T>
26where
27    T: Copy + Default + AddAssign + Sub<Output = T>,
28{
29    /// Build a Fenwick Tree from an initial slice of values.
30    ///
31    /// Runs in `O(n log n)`. For an `O(n)` build, build an empty tree
32    /// and call [`Self::update`] for each element.
33    pub fn new(nums: &[T]) -> Self {
34        let size = nums.len();
35        let tree = vec![T::default(); size + 1];
36
37        let mut ft = Self { size, tree };
38
39        for (i, &val) in nums.iter().enumerate() {
40            ft.update(i, val);
41        }
42        ft
43    }
44
45    /// Add `delta` to the element at `index`.
46    ///
47    /// Panics if `index >= self.size`.
48    pub fn update(&mut self, index: usize, delta: T) {
49        if index >= self.size {
50            panic!("Index out of bounds");
51        }
52
53        let mut i = index + 1;
54        loop {
55            if i > self.size { break }
56            self.tree[i] += delta;
57            i += i & (!i + 1); // equivalent to i & -i
58        }
59    }
60
61    /// Compute the prefix sum of elements `[0, index]` (inclusive).
62    ///
63    /// Panics if `index >= self.size`.
64    pub fn pref(&self, index: usize) -> T {
65        if index >= self.size {
66            panic!("Index out of bounds");
67        }
68        let mut result = T::default();
69        let mut i = index + 1;
70        loop {
71            if i <= 0 { break }
72            result += self.tree[i];
73            i -= i & (!i + 1); // equivalent to i & -i
74        }
75        result
76    }
77
78    /// Compute the sum of elements in the half-open range `[left, right]`.
79    ///
80    /// Panics if `right < left` or `right >= self.size`.
81    pub fn sum_range(&self, left: usize, right: usize) -> T {
82        if right < left || right >= self.size {
83            panic!("Invalid range");
84        }
85        match left == 0 {
86            true => self.pref(right),
87            false => self.pref(right) - self.pref(left - 1)
88        }
89    }
90}
fenwick/lib.rs

fenwick/
lib.rs
