nekolib/ds/
vec_segtree.rs

1//! `Vec` ベースのセグ木。
2
3use super::super::traits::binop;
4use super::super::traits::fold;
5use super::super::traits::fold_bisect;
6use super::super::traits::get_mut;
7use super::super::traits::set_value;
8use super::super::utils::buf_range;
9
10use std::convert::From;
11use std::fmt::{self, Debug};
12use std::iter::{IntoIterator, Iterator};
13use std::ops::{Deref, DerefMut, Index, Range, RangeBounds};
14
15use binop::Monoid;
16use buf_range::{bounds_within, check_bounds, check_bounds_range};
17use fold::Fold;
18use fold_bisect::{FoldBisect, FoldBisectRev};
19use get_mut::GetMut;
20use set_value::SetValue;
21
22/// `Vec` ベースのセグ木。
23///
24/// 非再帰実装かつ配列サイズを $2n$ とするセグ木。
25/// モノイドを対象として要素の取得・更新および任意区間のモノイド積を処理する。
26/// 加えて、モノイド積を引数とする述語に対して、それが真となる区間の境界を求められる。
27///
28/// # Complexity
29/// |演算|時間計算量|
30/// |---|---|
31/// |`new`, `from`|$\\Theta(n)$|
32/// |`index`|$\\Theta(1)$|
33/// |`set_value`|$\\Theta(\\log(n))$|
34/// |区間 $[l, r)$ の `fold`|$\\Theta(\\log(r-l))$|
35/// |`fold_bisect`, `fold_bisect_rev`|$O(\\log(n))$|
36///
37/// # Examples
38/// ```
39/// use nekolib::ds::VecSegtree;
40/// use nekolib::traits::{Fold, FoldBisect, FoldBisectRev, GetMut};
41/// use nekolib::utils::OpAdd;
42///
43/// let mut vs: VecSegtree<OpAdd<i32>> = vec![2, 4, 7, 3, 5].into();
44/// assert_eq!(vs.fold(1..3), 11);
45/// assert_eq!(vs.fold(..), 21);
46///
47/// *vs.get_mut(2).unwrap() = 1; // [2, 4, 1, 3, 5]
48/// assert_eq!(vs.fold(1..3), 5);
49/// assert_eq!(vs.fold(1..=3), 8);
50/// assert_eq!(vs.fold_bisect(1, |&x| x < 4), (1_usize, 0));
51/// assert_eq!(vs.fold_bisect(1, |&x| x <= 4), (2_usize, 4));
52/// assert_eq!(vs.fold_bisect(1, |&x| x < 13), (4_usize, 8));
53/// assert_eq!(vs.fold_bisect(1, |&x| x <= 13), (5_usize, 13));
54///
55/// assert_eq!(vs.fold(..), 15);
56/// assert_eq!(vs.fold_bisect_rev(5, |&x| x <= 0), (5_usize, 0));
57/// assert_eq!(vs.fold_bisect_rev(5, |&x| x < 15), (1_usize, 13));
58/// assert_eq!(vs.fold_bisect_rev(5, |&x| x <= 15), (0_usize, 15));
59///
60/// let l = 1;
61/// let pred = |&x: &i32| x <= 12;
62/// let (r, x) = vs.fold_bisect(l, pred);
63/// assert_eq!(vs.fold(l..r), x);
64/// assert!(pred(&x));
65/// assert!(r == vs.len() || !pred(&vs.fold(l..r + 1)));
66///
67/// let r = 5;
68/// let pred = |&x: &i32| x <= 12;
69/// let (l, x) = vs.fold_bisect_rev(r, pred);
70/// assert_eq!(vs.fold(l..r), x);
71/// assert!(pred(&x));
72/// assert!(l == 0 || !pred(&vs.fold(l - 1..r)));
73/// ```
74#[derive(Clone)]
75pub struct VecSegtree<M>
76where
77    M: Monoid,
78    M::Set: Clone,
79{
80    buf: Vec<M::Set>,
81    len: usize,
82    monoid: M,
83}
84
85impl<M> VecSegtree<M>
86where
87    M: Monoid,
88    M::Set: Clone,
89{
90    #[must_use]
91    pub fn new(len: usize) -> Self
92    where
93        M: Default,
94    {
95        let monoid = M::default();
96        Self { len, buf: vec![monoid.id(); len + len], monoid }
97    }
98
99    pub fn is_empty(&self) -> bool { self.len == 0 }
100
101    pub fn len(&self) -> usize { self.len }
102
103    fn nodes(&self, l: usize, r: usize) -> Vec<usize> {
104        let mut l = self.len + l;
105        let mut r = self.len + r;
106        let mut vl = vec![];
107        let mut vr = vec![];
108        while l < r {
109            if l & 1 == 1 {
110                vl.push(l);
111                l += 1;
112            }
113            if r & 1 == 1 {
114                r -= 1;
115                vr.push(r);
116            }
117            l >>= 1;
118            r >>= 1;
119        }
120        vr.reverse();
121        vl.append(&mut vr);
122        vl
123    }
124
125    fn nodes_rev(&self, l: usize, r: usize) -> Vec<usize> {
126        self.nodes(l, r).into_iter().rev().collect()
127    }
128}
129
130impl<M, B> Fold<B> for VecSegtree<M>
131where
132    M: Monoid,
133    M::Set: Clone,
134    B: RangeBounds<usize>,
135{
136    type Output = M;
137    fn fold(&self, b: B) -> M::Set {
138        let Range { start, end } = bounds_within(b, self.len);
139        let mut il = self.len + start;
140        let mut ir = self.len + end;
141        let mut res_l = self.monoid.id();
142        let mut res_r = self.monoid.id();
143        while il < ir {
144            if il & 1 == 1 {
145                res_l = self.monoid.op(res_l, self.buf[il].clone());
146                il += 1;
147            }
148            if ir & 1 == 1 {
149                ir -= 1;
150                res_r = self.monoid.op(self.buf[ir].clone(), res_r);
151            }
152            il >>= 1;
153            ir >>= 1;
154        }
155        self.monoid.op(res_l, res_r)
156    }
157}
158
159impl<M> SetValue<usize> for VecSegtree<M>
160where
161    M: Monoid,
162    M::Set: Clone,
163{
164    type Input = M::Set;
165    fn set_value(&mut self, i: usize, x: Self::Input) {
166        check_bounds(i, self.len);
167        *self.get_mut(i).unwrap() = x;
168    }
169}
170
171#[doc(hidden)]
172pub struct GetMutIndex<'a, M>
173where
174    M: Monoid,
175    M::Set: Clone,
176{
177    tree: &'a mut VecSegtree<M>,
178    index: usize,
179}
180
181impl<'a, M: 'a> GetMut<'a> for VecSegtree<M>
182where
183    M: Monoid,
184    M::Set: Clone,
185{
186    type Output = GetMutIndex<'a, M>;
187    fn get_mut(&'a mut self, index: usize) -> Option<GetMutIndex<'a, M>> {
188        if index < self.len {
189            Some(GetMutIndex { tree: self, index })
190        } else {
191            None
192        }
193    }
194}
195
196impl<M> Drop for GetMutIndex<'_, M>
197where
198    M: Monoid,
199    M::Set: Clone,
200{
201    fn drop(&mut self) {
202        let mut i = self.tree.len + self.index;
203        while i > 1 {
204            i >>= 1;
205            self.tree.buf[i] = self.tree.monoid.op(
206                self.tree.buf[i << 1].clone(),
207                self.tree.buf[i << 1 | 1].clone(),
208            );
209        }
210    }
211}
212
213impl<M> Deref for GetMutIndex<'_, M>
214where
215    M: Monoid,
216    M::Set: Clone,
217{
218    type Target = M::Set;
219    fn deref(&self) -> &Self::Target {
220        &self.tree.buf[self.tree.len + self.index]
221    }
222}
223
224impl<M> DerefMut for GetMutIndex<'_, M>
225where
226    M: Monoid,
227    M::Set: Clone,
228{
229    fn deref_mut(&mut self) -> &mut Self::Target {
230        &mut self.tree.buf[self.tree.len + self.index]
231    }
232}
233
234impl<M> Default for VecSegtree<M>
235where
236    M: Monoid + Default,
237    M::Set: Clone,
238{
239    fn default() -> Self { Self { buf: vec![], len: 0, monoid: M::default() } }
240}
241
242impl<M> From<Vec<M::Set>> for VecSegtree<M>
243where
244    M: Monoid + Default,
245    M::Set: Clone,
246{
247    fn from(v: Vec<M::Set>) -> Self { Self::from((v, M::default())) }
248}
249
250impl<M> From<(Vec<M::Set>, M)> for VecSegtree<M>
251where
252    M: Monoid,
253    M::Set: Clone,
254{
255    fn from((mut v, monoid): (Vec<M::Set>, M)) -> Self {
256        let len = v.len();
257        let mut buf = vec![monoid.id(); len];
258        buf.append(&mut v);
259        for i in (0..len).rev() {
260            buf[i] = monoid.op(buf[i << 1].clone(), buf[i << 1 | 1].clone());
261        }
262        Self { buf, len, monoid }
263    }
264}
265
266impl<M> Index<usize> for VecSegtree<M>
267where
268    M: Monoid,
269    M::Set: Clone,
270{
271    type Output = M::Set;
272    fn index(&self, i: usize) -> &Self::Output {
273        check_bounds(i, self.len);
274        &self.buf[i + self.len]
275    }
276}
277
278impl<M> From<VecSegtree<M>> for Vec<M::Set>
279where
280    M: Monoid,
281    M::Set: Clone,
282{
283    fn from(v: VecSegtree<M>) -> Self {
284        v.buf.into_iter().skip(v.len).collect()
285    }
286}
287
288impl<M> Debug for VecSegtree<M>
289where
290    M: Monoid,
291    M::Set: Clone + Debug,
292{
293    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
294        f.debug_list().entries(self.buf[self.len..].iter()).finish()
295    }
296}
297
298impl<M> FoldBisect for VecSegtree<M>
299where
300    M: Monoid,
301    M::Set: Clone,
302{
303    fn fold_bisect<F>(&self, l: usize, pred: F) -> (usize, M::Set)
304    where
305        F: Fn(&M::Set) -> bool,
306    {
307        check_bounds_range(l, 0..=self.len);
308
309        let mut x = self.monoid.id();
310        assert!(pred(&x), "`pred(id)` must hold");
311        match self.fold(l..) {
312            x if pred(&x) => return (self.len, x),
313            _ => {}
314        }
315
316        for v in self.nodes(l, self.len) {
317            let tmp = self.monoid.op(x.clone(), self.buf[v].clone());
318            if pred(&tmp) {
319                x = tmp;
320                continue;
321            }
322            let mut v = v;
323            while v < self.len {
324                v <<= 1;
325                let tmp = self.monoid.op(x.clone(), self.buf[v].clone());
326                if pred(&tmp) {
327                    x = tmp;
328                    v += 1;
329                }
330            }
331            return (v - self.len, x);
332        }
333        unreachable!();
334    }
335}
336
337impl<M> FoldBisectRev for VecSegtree<M>
338where
339    M: Monoid,
340    M::Set: Clone,
341{
342    fn fold_bisect_rev<F>(&self, r: usize, pred: F) -> (usize, M::Set)
343    where
344        F: Fn(&M::Set) -> bool,
345    {
346        check_bounds_range(r, 0..=self.len);
347
348        let mut x = self.monoid.id();
349        assert!(pred(&x), "`pred(id)` must hold");
350        match self.fold(..r) {
351            x if pred(&x) => return (0, x),
352            _ => {}
353        }
354
355        for v in self.nodes_rev(0, r) {
356            let tmp = self.monoid.op(self.buf[v].clone(), x.clone());
357            if pred(&tmp) {
358                x = tmp;
359                continue;
360            }
361            let mut v = v;
362            while v < self.len {
363                v = v << 1 | 1;
364                let tmp = self.monoid.op(self.buf[v].clone(), x.clone());
365                if pred(&tmp) {
366                    x = tmp;
367                    v -= 1;
368                }
369            }
370            return (v - self.len + 1, x);
371        }
372        unreachable!();
373    }
374}
nekolib/ds/vec_segtree.rs

nekolib/ds/
vec_segtree.rs
