🦀 Functional Rust

367: Fenwick Tree — Efficient Prefix Sums

Difficulty: 3 Level: Advanced O(log n) prefix sum queries and point updates in half the memory of a segment tree, powered by a single bit trick.

The Problem This Solves

You maintain an array of values that updates frequently. You often need the sum of all elements from index 0 to some index i (a "prefix sum"). A plain prefix sum array gives O(1) queries but O(n) updates — every update forces a full rebuild. A segment tree gives O(log n) for both but requires 4n nodes and two recursive functions. A Fenwick tree (Binary Indexed Tree, BIT) achieves O(log n) for both prefix sum queries and point updates using just n+1 storage and two loops — no recursion, no tree struct, no child pointers. The implementation is strikingly compact: 10-15 lines of code. It has excellent cache performance because everything is a contiguous array. The limitation: Fenwick trees are specialized for prefix sums (and by extension, range sums via prefix(r) - prefix(l-1)). They don't generalize to range minimum/maximum queries without significant complexity. For those, reach for a segment tree instead.

The Intuition

The bit trick that makes it work: `i & (-i)` isolates the lowest set bit of i. This determines the "responsibility range" of each array cell — how many elements it stores the sum of. To update index i: add the delta to tree[i], then jump to `i += i & (-i)` and repeat until you leave the array. You're walking up the chain of nodes that "cover" index i. To query prefix sum [1..i]: add tree[i], then jump to `i -= i & (-i)` and repeat until i = 0. You're collecting non-overlapping partial sums that together cover [1..i]. Think of it as a compressed partial sum table where each cell covers a power-of-two-sized range determined by its position's trailing zeros.

How It Works in Rust

struct Fenwick {
 tree: Vec<i64>, // 1-indexed; tree[0] unused
}

impl Fenwick {
 fn new(n: usize) -> Self {
     Fenwick { tree: vec![0; n + 1] }
 }

 // Build from existing data in O(n)
 fn from_slice(data: &[i64]) -> Self {
     let mut fw = Fenwick::new(data.len());
     for (i, &v) in data.iter().enumerate() {
         fw.add(i + 1, v); // Fenwick is 1-indexed
     }
     fw
 }

 // Add delta to index i (1-indexed) — O(log n)
 fn add(&mut self, mut i: usize, delta: i64) {
     while i < self.tree.len() {
         self.tree[i] += delta;
         i += i & i.wrapping_neg(); // i += lowest set bit of i
     }
 }

 // Prefix sum [1..=i] (1-indexed) — O(log n)
 fn prefix_sum(&self, mut i: usize) -> i64 {
     let mut sum = 0;
     while i > 0 {
         sum += self.tree[i];
         i -= i & i.wrapping_neg(); // i -= lowest set bit of i
     }
     sum
 }

 // Range sum [l..=r] (1-indexed) — O(log n)
 fn range_sum(&self, l: usize, r: usize) -> i64 {
     self.prefix_sum(r) - if l > 1 { self.prefix_sum(l - 1) } else { 0 }
 }

 // Point value at index i — O(log n) via range_sum
 fn get(&self, i: usize) -> i64 {
     self.range_sum(i, i)
 }
}

// Usage
let data = vec![3i64, 2, -1, 6, 5, 4, -3, 3, 7, 2];
let mut fw = Fenwick::from_slice(&data);

println!("{}", fw.prefix_sum(5)); // sum of first 5: 3+2-1+6+5 = 15
println!("{}", fw.range_sum(3, 7)); // sum of indices 3..=7

fw.add(4, 1); // add 1 to index 4 (now 6+1=7)
println!("{}", fw.prefix_sum(5)); // 16 after update

What This Unlocks

Key Differences

ConceptOCamlRust
Prefix sum (static)O(n) precompute + O(1) querysame, use prefix sum array
Prefix sum (dynamic)O(n) scan per queryO(log n) Fenwick
Point updateO(1)O(log n) propagation
SpaceO(n)O(n) — half of segment tree
Code complexitymoderatevery compact (2 loops)
Generalizes to range min/max?with effortno — use segment tree instead
struct FenwickTree {
    tree: Vec<i64>,
    n: usize,
}

impl FenwickTree {
    fn new(n: usize) -> Self {
        Self { tree: vec![0; n+1], n }
    }

    fn from_slice(arr: &[i64]) -> Self {
        let mut ft = Self::new(arr.len());
        for (i, &v) in arr.iter().enumerate() { ft.update(i+1, v); }
        ft
    }

    fn update(&mut self, mut i: usize, delta: i64) {
        while i <= self.n {
            self.tree[i] += delta;
            i += i & i.wrapping_neg();
        }
    }

    fn prefix_sum(&self, mut i: usize) -> i64 {
        let mut sum = 0;
        while i > 0 {
            sum += self.tree[i];
            i -= i & i.wrapping_neg();
        }
        sum
    }

    fn range_sum(&self, l: usize, r: usize) -> i64 {
        self.prefix_sum(r) - if l > 1 { self.prefix_sum(l-1) } else { 0 }
    }

    fn point_query(&self, i: usize) -> i64 {
        self.range_sum(i, i)
    }
}

fn main() {
    let arr: Vec<i64> = vec![1,2,3,4,5,6,7,8];
    let mut ft = FenwickTree::from_slice(&arr);
    println!("Prefix(4) = {}", ft.prefix_sum(4));    // 1+2+3+4=10
    println!("Range(2,5) = {}", ft.range_sum(2,5));  // 2+3+4+5=14
    ft.update(3, 10); // add 10 to position 3
    println!("After update(3,+10), Range(2,5) = {}", ft.range_sum(2,5)); // 2+13+4+5=24
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test] fn prefix_sums() {
        let ft = FenwickTree::from_slice(&[1,2,3,4,5]);
        assert_eq!(ft.prefix_sum(3), 6);
        assert_eq!(ft.prefix_sum(5), 15);
    }
    #[test] fn range_query() {
        let ft = FenwickTree::from_slice(&[1,2,3,4,5]);
        assert_eq!(ft.range_sum(2,4), 9);
    }
    #[test] fn point_update() {
        let mut ft = FenwickTree::from_slice(&[1,2,3,4,5]);
        ft.update(3, 7);
        assert_eq!(ft.point_query(3), 10);
        assert_eq!(ft.prefix_sum(5), 22);
    }
}
(* OCaml: Fenwick tree *)

let tree = Array.make 17 0  (* 1-indexed, size n+1 *)
let n = 16

let update i v =
  let i = ref i in
  while !i <= n do tree.(!i) <- tree.(!i) + v; i := !i + (!i land (- !i)) done

let prefix_sum i =
  let i = ref i and s = ref 0 in
  while !i > 0 do s := !s + tree.(!i); i := !i - (!i land (- !i)) done;
  !s

let range_sum l r = prefix_sum r - prefix_sum (l-1)

let () =
  List.iteri (fun i v -> update (i+1) v) [1;2;3;4;5;6;7;8];
  Printf.printf "Prefix(4) = %d\n" (prefix_sum 4);  (* 1+2+3+4=10 *)
  Printf.printf "Range(2,5) = %d\n" (range_sum 2 5) (* 2+3+4+5=14 *)