ac_lib/structure/
segtree.rs

1pub struct SegmentTree {
2    size: usize,
3    tree: Vec<i64>,
4    n: usize,
5}
6
7impl SegmentTree {
8    pub fn new(arr: &[i64]) -> Self {
9        let n = arr.len();
10        let size = n.next_power_of_two();
11        let mut tree = vec![0; size * 2];
12
13        tree[size..(n + size)].copy_from_slice(&arr[..n]);
14
15        for i in (1..size).rev() {
16            tree[i] = tree[i * 2] + tree[i * 2 + 1];
17        }
18
19        SegmentTree { size, tree, n }
20    }
21
22    pub fn update(&mut self, idx: usize, value: i64) {
23        assert!(idx < self.n, "Index out of bounds");
24
25        let mut pos = self.size + idx;
26        self.tree[pos] = value;
27
28        while pos > 1 {
29            pos /= 2;
30            self.tree[pos] = self.tree[pos * 2] + self.tree[pos * 2 + 1];
31        }
32    }
33
34    pub fn query(&self, l: usize, r: usize) -> i64 {
35        assert!(l <= r && r <= self.n, "Invalid range");
36
37        let mut left = self.size + l;
38        let mut right = self.size + r;
39        let mut sum = 0;
40
41        while left < right {
42            if left % 2 == 1 {
43                sum += self.tree[left];
44                left += 1;
45            }
46            if right % 2 == 1 {
47                right -= 1;
48                sum += self.tree[right];
49            }
50            left /= 2;
51            right /= 2;
52        }
53
54        sum
55    }
56
57    pub fn get(&self, idx: usize) -> i64 {
58        assert!(idx < self.n, "Index out of bounds");
59        self.tree[self.size + idx]
60    }
61}
ac_lib/structure/segtree.rs

ac_lib/structure/
segtree.rs
