ac_lib/structure/
fenwick.rs

1pub struct FenwickTree {
2    size: usize,
3    tree: Vec<i64>,
4}
5
6impl FenwickTree {
7    pub fn new(size: usize) -> Self {
8        FenwickTree {
9            size,
10            tree: vec![0; size + 1],
11        }
12    }
13
14    pub fn from_vec(arr: &[i64]) -> Self {
15        let mut ft = Self::new(arr.len());
16        for (i, &val) in arr.iter().enumerate() {
17            ft.add(i, val);
18        }
19        ft
20    }
21
22    pub fn add(&mut self, index: usize, value: i64) {
23        assert!(index < self.size, "Index out of bounds");
24
25        let mut idx = index + 1;
26        while idx <= self.size {
27            self.tree[idx] += value;
28            idx += idx & (!idx + 1);
29        }
30    }
31
32    pub fn sum(&self, index: usize) -> i64 {
33        if index >= self.size {
34            return self.sum(self.size - 1);
35        }
36
37        let mut sum = 0;
38        let mut idx = index + 1;
39        while idx > 0 {
40            sum += self.tree[idx];
41            idx &= idx - 1;
42        }
43        sum
44    }
45
46    pub fn range_sum(&self, left: usize, right: usize) -> i64 {
47        assert!(left <= right && right < self.size, "Invalid range");
48
49        if left == 0 {
50            self.sum(right)
51        } else {
52            self.sum(right) - self.sum(left - 1)
53        }
54    }
55
56    pub fn get(&self, index: usize) -> i64 {
57        self.range_sum(index, index)
58    }
59
60    pub fn set(&mut self, index: usize, value: i64) {
61        let current = self.get(index);
62        self.add(index, value - current);
63    }
64}
ac_lib/structure/fenwick.rs

ac_lib/structure/
fenwick.rs
