competitive/data_structure/
range_ap_add.rs

1#[derive(Debug, Clone)]
2pub struct RangeArithmeticProgressionAdd {
3    pub dd: Vec<i64>,
4}
5impl RangeArithmeticProgressionAdd {
6    pub fn new(n: usize) -> Self {
7        Self { dd: vec![0; n] }
8    }
9    /// add a, a+d, ..., a+(k-1)d into [l, l + k)
10    pub fn update(&mut self, l: usize, k: usize, a: i64, d: i64) {
11        if let Some(e) = self.dd.get_mut(l) {
12            *e += a;
13        }
14        if let Some(e) = self.dd.get_mut(l + 1) {
15            *e += d - a;
16        }
17        if let Some(e) = self.dd.get_mut(l + k) {
18            *e += -a - k as i64 * d;
19        }
20        if let Some(e) = self.dd.get_mut(l + k + 1) {
21            *e += a + (k as i64 - 1) * d;
22        }
23    }
24    /// add a, a+d, ..., a+(k-1)d into [l, l + k)
25    pub fn update_isize(&mut self, l: isize, k: usize, a: i64, d: i64) {
26        if l < 0 {
27            let r = l + k as isize;
28            if r > 0 {
29                self.update(0, r as usize, a - l as i64 * d, d);
30            }
31        } else {
32            self.update(l as usize, k, a, d);
33        }
34    }
35    pub fn build_inplace(&mut self) {
36        for _ in 0..2 {
37            for i in 0..self.dd.len() - 1 {
38                self.dd[i + 1] += self.dd[i];
39            }
40        }
41    }
42}
43
44#[test]
45fn test_range_ap_add() {
46    use crate::tools::{NotEmptySegment as Nes, Xorshift};
47    const N: usize = 1000;
48    const Q: usize = 10000;
49    const A: i64 = 1_000_000_000;
50    let mut rng = Xorshift::new();
51    let mut v = vec![0i64; N];
52    let mut ap = RangeArithmeticProgressionAdd::new(N);
53    for ((l, r), a, d) in rng.random_iter((Nes(N), -A..=A, -A..=A)).take(Q) {
54        for (i, v) in v[l..r].iter_mut().enumerate() {
55            *v += a + i as i64 * d;
56        }
57        ap.update(l, r - l, a, d);
58    }
59    ap.build_inplace();
60    assert_eq!(ap.dd, v);
61}
competitive/data_structure/range_ap_add.rs

competitive/data_structure/
range_ap_add.rs
