cp_library_rs/data_structure/
dual_segment_tree.rs

1//! 双対セグメント木:区間加算・一点取得
2
3use std::fmt::{self, Debug};
4use std::ops::{
5    Bound::{Excluded, Included, Unbounded},
6    RangeBounds,
7};
8
9use crate::algebraic_structure::commutative::CommutativeMonoid;
10
11/// 双対セグ木
12/// - 区間への作用
13/// - 一点の取得
14///
15/// を行うセグメント木
16pub struct DualSegmentTree<M: CommutativeMonoid> {
17    pub size: usize,
18    offset: usize,
19    data: Vec<M::Val>,
20}
21
22impl<M: CommutativeMonoid> DualSegmentTree<M> {
23    #[inline]
24    fn parse_range<R: RangeBounds<usize>>(&self, range: &R) -> Option<(usize, usize)> {
25        let start = match range.start_bound() {
26            Unbounded => 0,
27            Excluded(&v) => v + 1,
28            Included(&v) => v,
29        };
30        let end = match range.end_bound() {
31            Unbounded => self.size,
32            Excluded(&v) => v,
33            Included(&v) => v + 1,
34        };
35        if start <= end && end <= self.size {
36            Some((start, end))
37        } else {
38            None
39        }
40    }
41
42    /// 双対セグメント木を初期化する
43    pub fn new(n: usize) -> Self {
44        let offset = n;
45
46        Self {
47            size: n,
48            offset,
49            data: vec![M::e(); offset << 1],
50        }
51    }
52
53    /// 配列から双対セグメント木を構築する
54    pub fn from_vec(arr: Vec<M::Val>) -> Self {
55        let offset = arr.len();
56        let mut seg = Self::new(offset);
57        for (i, v) in arr.into_iter().enumerate() {
58            seg.data[offset + i] = v;
59        }
60        seg
61    }
62
63    /// 区間更新:
64    /// - 区間`range`を`x`で更新する
65    pub fn apply_range<R: RangeBounds<usize> + Debug>(&mut self, range: R, x: M::Val) {
66        let Some((start, end)) = self.parse_range(&range) else {
67            panic!("The given range is wrong: {:?}", range);
68        };
69        // 値の更新
70        let mut l = self.offset + start;
71        let mut r = self.offset + end;
72
73        while l < r {
74            if l & 1 == 1 {
75                let tmp = M::op(&self.data[l], &x);
76                self.data[l] = tmp;
77                l += 1;
78            }
79            if r & 1 == 1 {
80                r -= 1;
81                let tmp = M::op(&self.data[r], &x);
82                self.data[r] = tmp;
83            }
84            l >>= 1;
85            r >>= 1;
86        }
87    }
88
89    /// 一点取得
90    pub fn get_point(&self, index: usize) -> M::Val {
91        let mut i = index + self.offset;
92        let mut res = self.data[i].clone();
93        while i > 1 {
94            i >>= 1;
95            let tmp = M::op(&self.data[i], &res);
96            res = tmp;
97        }
98        res
99    }
100}
101
102impl<M> Debug for DualSegmentTree<M>
103where
104    M: CommutativeMonoid,
105    M::Val: Debug,
106{
107    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
108        write!(f, "DualSegmentTree {{ [").ok();
109        for i in 0..self.size {
110            if i + 1 < self.size {
111                write!(f, "{:?}, ", self.get_point(i)).ok();
112            } else {
113                write!(f, "{:?}", self.get_point(i)).ok();
114            }
115        }
116        write!(f, "] }}")
117    }
118}