haar_lib/ds/
lazy_segtree_coeff.rs

1//! 係数乗算付き区間加算区間総和遅延セグ木
2
3use crate::misc::range::range_bounds_to_range;
4use crate::num::one_zero::Zero;
5use crate::trait_alias;
6use std::cell::Cell;
7use std::ops::{Add, Mul, RangeBounds};
8
9trait_alias!(
10    /// [`LazySegtreeCoeff<T>`]が扱える型
11    Elem: Copy + Zero + Add<Output = Self> + Mul<Output = Self> + PartialEq
12);
13
14/// 係数乗算付き区間加算区間総和遅延セグ木
15pub struct LazySegtreeCoeff<T> {
16    size: usize,
17    original_size: usize,
18    data: Vec<Cell<T>>,
19    lazy: Vec<Cell<T>>,
20    coeff: Vec<T>,
21}
22
23impl<T: Elem> LazySegtreeCoeff<T> {
24    /// ‍係数`coefficients`を設定した[`LazySegtreeCoeff`]を構築する。
25    pub fn new(n: usize, coefficients: Vec<T>) -> Self {
26        let size = n.next_power_of_two() * 2;
27
28        let mut coeff = vec![T::zero(); size];
29
30        for i in 0..coefficients.len() {
31            coeff[i + size / 2] = coefficients[i];
32        }
33        for i in (1..size / 2).rev() {
34            coeff[i] = coeff[i << 1] + coeff[(i << 1) | 1];
35        }
36
37        Self {
38            size,
39            original_size: n,
40            data: vec![Cell::new(T::zero()); size],
41            lazy: vec![Cell::new(T::zero()); size],
42            coeff,
43        }
44    }
45
46    /// `self.fold(i..i+1) = value[i]`となるように割り当てる。
47    pub fn set_vec(&mut self, value: Vec<T>) {
48        self.data = vec![Cell::new(T::zero()); self.size];
49        self.lazy = vec![Cell::new(T::zero()); self.size];
50
51        for (i, x) in value.into_iter().enumerate() {
52            self.data[self.size / 2 + i].set(x);
53        }
54        for i in (1..self.size / 2).rev() {
55            self.data[i].set(self.data[i << 1].get() + self.data[(i << 1) | 1].get());
56        }
57    }
58
59    fn propagate(&self, i: usize) {
60        if self.lazy[i].get() != T::zero() {
61            if i < self.size / 2 {
62                self.lazy[i << 1].set(self.lazy[i].get() + self.lazy[i << 1].get());
63                self.lazy[(i << 1) | 1].set(self.lazy[i].get() + self.lazy[(i << 1) | 1].get());
64            }
65            self.data[i].set(self.data[i].get() + self.lazy[i].get() * self.coeff[i]);
66            self.lazy[i].set(T::zero());
67        }
68    }
69
70    fn update_internal(&mut self, i: usize, l: usize, r: usize, s: usize, t: usize, value: T) -> T {
71        self.propagate(i);
72        if r <= s || t <= l {
73            return self.data[i].get();
74        }
75        if s <= l && r <= t {
76            self.lazy[i].set(self.lazy[i].get() + value);
77            self.propagate(i);
78            return self.data[i].get();
79        }
80
81        let t = self.update_internal(i << 1, l, (l + r) / 2, s, t, value)
82            + self.update_internal((i << 1) | 1, (l + r) / 2, r, s, t, value);
83
84        self.data[i].replace(t)
85    }
86
87    fn get_internal(&self, i: usize, l: usize, r: usize, x: usize, y: usize) -> T {
88        self.propagate(i);
89        if r <= x || y <= l {
90            return T::zero();
91        }
92        if x <= l && r <= y {
93            return self.data[i].get();
94        }
95        self.get_internal(i << 1, l, (l + r) / 2, x, y)
96            + self.get_internal((i << 1) | 1, (l + r) / 2, r, x, y)
97    }
98
99    /// 範囲`range`に値`value`を加算する。
100    pub fn update(&mut self, range: impl RangeBounds<usize>, value: T) {
101        let (start, end) = range_bounds_to_range(range, 0, self.original_size);
102        self.update_internal(1, 0, self.size / 2, start, end, value);
103    }
104
105    /// 範囲`range`で総和を取る。
106    pub fn fold(&self, range: impl RangeBounds<usize>) -> T {
107        let (start, end) = range_bounds_to_range(range, 0, self.original_size);
108        self.get_internal(1, 0, self.size / 2, start, end)
109    }
110}
haar_lib/ds/lazy_segtree_coeff.rs

haar_lib/ds/
lazy_segtree_coeff.rs
