competitive_library/structure/
segment_tree.rs

1//! セグメントツリー
2use std::cmp::max;
3use std::cmp::min;
4
5pub trait Monoid {
6    type T: Clone;
7    fn identity_element() -> Self::T;
8    fn binary_operation(a: &Self::T, b: &Self::T) -> Self::T;
9}
10pub struct Min {}
11impl Monoid for Min {
12    type T = i64;
13    #[inline]
14    fn identity_element() -> Self::T {
15        i32::MAX as i64
16    }
17    #[inline]
18    fn binary_operation(a: &Self::T, b: &Self::T) -> Self::T {
19        min(*a, *b)
20    }
21}
22pub struct Max {}
23impl Monoid for Max {
24    type T = i64;
25    #[inline]
26    fn identity_element() -> Self::T {
27        i64::MIN
28    }
29    #[inline]
30    fn binary_operation(a: &Self::T, b: &Self::T) -> Self::T {
31        max(*a, *b)
32    }
33}
34
35pub struct SegmentTree<M>
36where
37    M: Monoid,
38{
39    n: usize,
40    size: usize,
41    log: usize,
42    data: Vec<M::T>,
43}
44impl<M: Monoid> SegmentTree<M> {
45    pub fn new(n: usize) -> SegmentTree<M> {
46        vec![M::identity_element(); n].into()
47    }
48}
49impl<M: Monoid> From<Vec<M::T>> for SegmentTree<M> {
50    fn from(v: Vec<M::T>) -> Self {
51        let n = v.len();
52        let log = (32 - (n as u32).saturating_sub(1).leading_zeros()) as usize;
53        let size = 1 << log;
54        let data = {
55            let mut data = vec![M::identity_element(); 2 * size];
56            data[size..(size + n)].clone_from_slice(&v);
57            data
58        };
59        {
60            let mut sg = SegmentTree { n, size, log, data };
61            (1..size).rev().for_each(|i| sg.update(i));
62            sg
63        }
64    }
65}
66impl<M: Monoid> SegmentTree<M> {
67    pub fn query(&self, mut l: usize, mut r: usize) -> M::T {
68        let (mut sml, mut smr) = (M::identity_element(), M::identity_element());
69        l += self.size;
70        r += self.size;
71
72        while l < r {
73            if l & 1 != 0 {
74                sml = M::binary_operation(&sml, unsafe { self.data.get_unchecked(l) });
75                l += 1;
76            }
77            if r & 1 != 0 {
78                r -= 1;
79                smr = M::binary_operation(unsafe { self.data.get_unchecked(r) }, &smr);
80            }
81            l >>= 1;
82            r >>= 1;
83        }
84
85        M::binary_operation(&sml, &smr)
86    }
87    fn update(&mut self, k: usize) {
88        *unsafe { self.data.get_unchecked_mut(k) } =
89            M::binary_operation(unsafe { self.data.get_unchecked(2 * k) }, unsafe {
90                self.data.get_unchecked(2 * k + 1)
91            });
92    }
93    pub fn set(&mut self, mut p: usize, x: M::T) {
94        p += self.size;
95        self.data[p] = x;
96        (1..=self.log).for_each(|i| self.update(p >> i));
97    }
98    pub fn get(&self, i: usize) -> M::T {
99        self.data[i].clone()
100    }
101    pub fn max_right<F>(&self, mut l: usize, f: F) -> usize
102    where
103        F: Fn(&M::T) -> bool,
104    {
105        assert!(l <= self.n);
106        assert!(f(&M::identity_element()));
107        if l == self.n {
108            return self.n;
109        }
110        l += self.size;
111        let mut sm = M::identity_element();
112        loop {
113            // do
114            while l % 2 == 0 {
115                l >>= 1;
116            }
117            if !f(&M::binary_operation(&sm, unsafe {
118                self.data.get_unchecked(l)
119            })) {
120                while l < self.size {
121                    l *= 2;
122                    let res = M::binary_operation(&sm, unsafe { self.data.get_unchecked(l) });
123                    if f(&res) {
124                        sm = res;
125                        l += 1;
126                    }
127                }
128                return l - self.size;
129            }
130            sm = M::binary_operation(&sm, unsafe { self.data.get_unchecked(l) });
131            l += 1;
132            // while
133
134            let l = l as isize;
135            if (l & -l) == l {
136                break;
137            }
138        }
139        self.n
140    }
141
142    pub fn min_left<F>(&self, mut r: usize, f: F) -> usize
143    where
144        F: Fn(&M::T) -> bool,
145    {
146        assert!(r <= self.n);
147        assert!(f(&M::identity_element()));
148        if r == 0 {
149            return 0;
150        }
151        r += self.size;
152        let mut sm = M::identity_element();
153        loop {
154            // do
155            r -= 1;
156            while r > 1 && r % 2 == 1 {
157                r >>= 1;
158            }
159            if !f(&M::binary_operation(
160                unsafe { self.data.get_unchecked(r) },
161                &sm,
162            )) {
163                while r < self.size {
164                    r = 2 * r + 1;
165                    let res = M::binary_operation(unsafe { self.data.get_unchecked(r) }, &sm);
166                    if f(&res) {
167                        sm = res;
168                        r -= 1;
169                    }
170                }
171                return r + 1 - self.size;
172            }
173            sm = M::binary_operation(unsafe { self.data.get_unchecked(r) }, &sm);
174            // while
175
176            let r = r as isize;
177            if (r & -r) == r {
178                break;
179            }
180        }
181
182        0
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    #[test]
190    fn test_practice2_sample() {
191        let a = vec![1, 2, 3, 2, 1];
192        let cxy = vec![(2, 1, 5), (3, 2, 3), (1, 3, 1), (2, 2, 4), (3, 1, 3)];
193
194        let mut st = SegmentTree::<Max>::from(a);
195        let mut ans = vec![];
196        for (c, x, y) in cxy {
197            if c == 1 {
198                st.set(x as usize - 1, y);
199            } else if c == 2 {
200                ans.push(st.query(x as usize - 1, y as usize));
201            } else {
202                ans.push(st.max_right(x as usize - 1, |&v| y > v) as i64 + 1);
203            }
204        }
205
206        assert_eq!(vec![3, 3, 2, 6], ans);
207    }
208}
competitive_library/structure/segment_tree.rs

competitive_library/structure/
segment_tree.rs
