nekolib/ds/
vec_act_segtree.rs

1//! `Vec` ベースの区間作用セグ木。
2
3use super::super::traits::act;
4use super::super::traits::action;
5use super::super::traits::binop;
6use super::super::traits::fold;
7use super::super::traits::fold_bisect;
8use super::super::traits::get_mut;
9use super::super::utils::buf_range;
10
11use std::cell::RefCell;
12use std::fmt::{self, Debug};
13use std::ops::{Deref, DerefMut, Range, RangeBounds};
14
15use act::Act;
16use action::MonoidAction;
17use binop::{Identity, Magma};
18use buf_range::{bounds_within, check_bounds_range};
19use fold::Fold;
20use fold_bisect::{FoldBisect, FoldBisectRev};
21use get_mut::GetMut;
22
23const WORD_SIZE: u32 = 0_usize.count_zeros();
24
25fn lcp(i: usize, j: usize) -> usize {
26    if i == 0 || j == 0 {
27        return 0;
28    }
29    if i == j {
30        return i;
31    }
32    let (i, j) = (i.min(j), i.max(j));
33    let iz = i.leading_zeros();
34    let i = i << iz;
35    let j = j << j.leading_zeros();
36    i >> iz.max(WORD_SIZE - (i ^ j).leading_zeros())
37}
38
39#[derive(Clone, Default)]
40pub struct VecActSegtree<A>
41where
42    A: MonoidAction,
43    <A::Operator as Magma>::Set: Clone,
44    <A::Operand as Magma>::Set: Clone,
45{
46    buf: RefCell<Vec<<A::Operand as Magma>::Set>>,
47    def: RefCell<Vec<<A::Operator as Magma>::Set>>,
48    len: usize,
49    action: A,
50}
51
52impl<A> VecActSegtree<A>
53where
54    A: MonoidAction,
55    <A::Operator as Magma>::Set: Clone,
56    <A::Operand as Magma>::Set: Clone,
57{
58    #[must_use]
59    pub fn new(len: usize) -> Self
60    where
61        A: Default,
62    {
63        let action = A::default();
64        Self {
65            len,
66            buf: RefCell::new(vec![action.operand().id(); len + len]),
67            def: RefCell::new(vec![action.operator().id(); len]),
68            action,
69        }
70    }
71
72    pub fn is_empty(&self) -> bool { self.len == 0 }
73    pub fn len(&self) -> usize { self.len }
74
75    fn arch_pair(&self, l: usize, r: usize) -> (Vec<usize>, Vec<usize>) {
76        let mut l = self.len + l;
77        let mut r = self.len + r;
78        let mut vl = vec![];
79        let mut vr = vec![];
80        while l < r {
81            if l & 1 == 1 {
82                vl.push(l);
83                l += 1;
84            }
85            if r & 1 == 1 {
86                r -= 1;
87                vr.push(r);
88            }
89            l >>= 1;
90            r >>= 1;
91        }
92        (vl, vr)
93    }
94
95    fn arch(&self, l: usize, r: usize) -> Vec<usize> {
96        let (mut vl, vr) = self.arch_pair(l, r);
97        vl.extend(vr.into_iter().rev());
98        vl
99    }
100
101    fn arch_rev(&self, l: usize, r: usize) -> Vec<usize> {
102        let (vl, mut vr) = self.arch_pair(l, r);
103        vr.extend(vl.into_iter().rev());
104        vr
105    }
106
107    fn build_range(&mut self, start: usize, end: usize) {
108        let mut buf = self.buf.borrow_mut();
109        let def = self.def.borrow();
110        let action = &self.action;
111        let operand = action.operand();
112        let id = action.operator().id();
113        for i in self.ancestors_upward(start, end).filter(|&i| def[i] == id) {
114            buf[i] = operand.op(buf[i << 1].clone(), buf[i << 1 | 1].clone());
115        }
116    }
117
118    fn apply(&self, i: usize, op: <A::Operator as Magma>::Set) {
119        let mut buf = self.buf.borrow_mut();
120        let mut def = self.def.borrow_mut();
121        buf[i] = self.action.act(buf[i].clone(), op.clone());
122        if i < self.len {
123            def[i] = self.action.operator().op(def[i].clone(), op);
124        }
125    }
126
127    fn force(&self, i: usize) {
128        let e = self.action.operator().id();
129        let d = {
130            let mut def = self.def.borrow_mut();
131            std::mem::replace(&mut def[i], e.clone())
132        };
133        if d != e {
134            self.apply(i << 1, d.clone());
135            self.apply(i << 1 | 1, d);
136        }
137    }
138
139    fn parent_root(&self, i: usize) -> usize {
140        let n = self.len;
141        if n.is_power_of_two() {
142            return 0;
143        }
144        let n2 = 2 * n;
145        let lsb = n2 & n2.wrapping_neg();
146        lcp(i, if i < n2 ^ lsb { n2 ^ lsb } else { n2 })
147    }
148
149    fn ancestors_downward(
150        &self,
151        start: usize,
152        end: usize,
153    ) -> impl Iterator<Item = usize> + DoubleEndedIterator {
154        self.ancestors_upward(start, end).rev()
155    }
156
157    fn ancestors_upward(
158        &self,
159        start: usize,
160        end: usize,
161    ) -> impl Iterator<Item = usize> + DoubleEndedIterator {
162        // start <= end
163
164        let mut res = vec![];
165        let l = self.len + start;
166        let r = self.len + end;
167        let pl = self.parent_root(l);
168        let pr = self.parent_root(r - 1);
169
170        let mut il = 1;
171        let mut ir = 1;
172        while l >> il != pl || (r - 1) >> ir != pr {
173            if l >> il != pl {
174                if l >> il << il != l {
175                    res.push(l >> il);
176                }
177                il += 1;
178            }
179            if r >> ir != pr {
180                if r >> ir << ir != r {
181                    res.push((r - 1) >> ir);
182                }
183                ir += 1;
184            }
185        }
186
187        res.dedup();
188        res.into_iter()
189    }
190
191    fn force_range(&self, l: usize, r: usize) {
192        for i in self.ancestors_downward(l, r) {
193            self.force(i);
194        }
195    }
196
197    fn force_all(&self) {
198        let mut buf = self.buf.borrow_mut();
199        let mut def = self.def.borrow_mut();
200        let operator = &self.action.operator();
201        let e = operator.id();
202        for i in 1..self.len {
203            let d = std::mem::replace(&mut def[i], e.clone());
204            for &j in &[i << 1, i << 1 | 1] {
205                if j < self.len {
206                    def[j] = operator.op(def[j].clone(), d.clone());
207                }
208                buf[j] = self.action.act(buf[j].clone(), d.clone());
209            }
210        }
211    }
212
213    #[cfg(test)]
214    pub fn unforced(&self) -> Vec<usize> {
215        let def = self.def.borrow();
216        let id = self.action.operator().id();
217        (0..self.len).filter(|&i| def[i] != id).collect()
218    }
219
220    #[cfg(test)]
221    pub fn ancestors_sorted(&self, start: usize, end: usize) -> Vec<usize> {
222        let mut res: Vec<_> = self.ancestors_upward(start, end).collect();
223        res.sort_unstable();
224        res
225    }
226}
227
228impl<A> From<Vec<<A::Operand as Magma>::Set>> for VecActSegtree<A>
229where
230    A: MonoidAction + Default,
231    <A::Operator as Magma>::Set: Clone,
232    <A::Operand as Magma>::Set: Clone,
233{
234    fn from(v: Vec<<A::Operand as Magma>::Set>) -> Self {
235        Self::from((v, A::default()))
236    }
237}
238
239impl<A> From<(Vec<<A::Operand as Magma>::Set>, A)> for VecActSegtree<A>
240where
241    A: MonoidAction,
242    <A::Operator as Magma>::Set: Clone,
243    <A::Operand as Magma>::Set: Clone,
244{
245    fn from((mut v, action): (Vec<<A::Operand as Magma>::Set>, A)) -> Self {
246        let len = v.len();
247        let mut buf = vec![action.operand().id(); len];
248        buf.append(&mut v);
249        for i in (0..len).rev() {
250            buf[i] = action
251                .operand()
252                .op(buf[i << 1].clone(), buf[i << 1 | 1].clone());
253        }
254        let buf = RefCell::new(buf);
255        let def = RefCell::new(vec![action.operator().id(); len]);
256        Self { buf, def, len, action }
257    }
258}
259
260impl<A> Debug for VecActSegtree<A>
261where
262    A: MonoidAction,
263    <A::Operator as Magma>::Set: Clone,
264    <A::Operand as Magma>::Set: Clone + Debug,
265{
266    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
267        self.force_all();
268        f.debug_list().entries(self.buf.borrow()[self.len..].iter()).finish()
269    }
270}
271
272impl<A> From<VecActSegtree<A>> for Vec<<A::Operand as Magma>::Set>
273where
274    A: MonoidAction,
275    <A::Operator as Magma>::Set: Clone,
276    <A::Operand as Magma>::Set: Clone,
277{
278    fn from(v: VecActSegtree<A>) -> Self {
279        v.force_all();
280        let mut res = v.buf.into_inner();
281        res.drain(..v.len);
282        res
283    }
284}
285
286impl<A, B> Fold<B> for VecActSegtree<A>
287where
288    A: MonoidAction,
289    <A::Operator as Magma>::Set: Clone,
290    <A::Operand as Magma>::Set: Clone,
291    B: RangeBounds<usize>,
292{
293    type Output = A::Operand;
294    fn fold(&self, b: B) -> <Self::Output as Magma>::Set {
295        let Range { start, end } = bounds_within(b, self.len);
296        let operand = self.action.operand();
297        if start >= end {
298            return operand.id();
299        }
300        self.force_range(start, end);
301        let mut res = operand.id();
302        let buf = self.buf.borrow();
303        for v in self.arch(start, end) {
304            res = operand.op(res, buf[v].clone());
305        }
306        res
307    }
308}
309
310impl<A, B> Act<B> for VecActSegtree<A>
311where
312    A: MonoidAction,
313    <A::Operator as Magma>::Set: Clone,
314    <A::Operand as Magma>::Set: Clone,
315    B: RangeBounds<usize>,
316{
317    type Action = A;
318    fn act(&mut self, b: B, op: <A::Operator as Magma>::Set) {
319        let Range { start, end } = bounds_within(b, self.len);
320        if start >= end {
321            return;
322        }
323        self.force_range(start, end);
324        for v in self.arch(start, end) {
325            self.apply(v, op.clone());
326        }
327        self.build_range(start, end);
328    }
329}
330
331impl<A> FoldBisect for VecActSegtree<A>
332where
333    A: MonoidAction,
334    <A::Operator as Magma>::Set: Clone,
335    <A::Operand as Magma>::Set: Clone,
336{
337    fn fold_bisect<F>(
338        &self,
339        l: usize,
340        pred: F,
341    ) -> (usize, <A::Operand as Magma>::Set)
342    where
343        F: Fn(&<A::Operand as Magma>::Set) -> bool,
344    {
345        check_bounds_range(l, 0..=self.len);
346
347        let operand = self.action.operand();
348        let mut x = operand.id();
349        assert!(pred(&x), "`pred(id)` must hold");
350        match self.fold(l..) {
351            x if pred(&x) => return (self.len, x),
352            _ => {}
353        }
354
355        self.force_range(l, self.len);
356        let buf = || self.buf.borrow();
357        for v in self.arch(l, self.len) {
358            let tmp = operand.op(x.clone(), buf()[v].clone());
359            if pred(&tmp) {
360                x = tmp;
361                continue;
362            }
363            let mut v = v;
364            while v < self.len {
365                self.force(v);
366                v <<= 1;
367                let tmp = operand.op(x.clone(), buf()[v].clone());
368                if pred(&tmp) {
369                    x = tmp;
370                    v += 1;
371                }
372            }
373            return (v - self.len, x);
374        }
375        unreachable!();
376    }
377}
378
379impl<A> FoldBisectRev for VecActSegtree<A>
380where
381    A: MonoidAction,
382    <A::Operator as Magma>::Set: Clone,
383    <A::Operand as Magma>::Set: Clone,
384{
385    fn fold_bisect_rev<F>(
386        &self,
387        r: usize,
388        pred: F,
389    ) -> (usize, <A::Operand as Magma>::Set)
390    where
391        F: Fn(&<A::Operand as Magma>::Set) -> bool,
392    {
393        check_bounds_range(r, 0..=self.len);
394
395        let operand = self.action.operand();
396        let mut x = operand.id();
397        assert!(pred(&x), "`pred(id)` must hold");
398        match self.fold(..r) {
399            x if pred(&x) => return (0, x),
400            _ => {}
401        }
402
403        self.force_range(0, r);
404        let buf = || self.buf.borrow();
405        for v in self.arch_rev(0, r) {
406            let tmp = operand.op(buf()[v].clone(), x.clone());
407            if pred(&tmp) {
408                x = tmp;
409                continue;
410            }
411            let mut v = v;
412            while v < self.len {
413                self.force(v);
414                v = v << 1 | 1;
415                let tmp = operand.op(buf()[v].clone(), x.clone());
416                if pred(&tmp) {
417                    x = tmp;
418                    v -= 1;
419                }
420            }
421            return (v - self.len + 1, x);
422        }
423        unreachable!();
424    }
425}
426
427#[doc(hidden)]
428pub struct GetMutIndex<'a, A>
429where
430    A: MonoidAction,
431    <A::Operator as Magma>::Set: Clone,
432    <A::Operand as Magma>::Set: Clone,
433{
434    tree: &'a mut VecActSegtree<A>,
435    index: usize,
436    elt: <A::Operand as Magma>::Set,
437}
438
439impl<'a, A: 'a> GetMut<'a> for VecActSegtree<A>
440where
441    A: MonoidAction,
442    <A::Operator as Magma>::Set: Clone,
443    <A::Operand as Magma>::Set: Clone,
444{
445    type Output = GetMutIndex<'a, A>;
446    fn get_mut(&'a mut self, index: usize) -> Option<GetMutIndex<'a, A>> {
447        if index >= self.len {
448            return None;
449        }
450
451        self.force_range(index, index + 1);
452        let i = self.len + index;
453        let e = self.action.operand().id();
454        let elt = std::mem::replace(&mut self.buf.borrow_mut()[i], e);
455        Some(GetMutIndex { tree: self, index, elt })
456    }
457}
458
459impl<A> Drop for GetMutIndex<'_, A>
460where
461    A: MonoidAction,
462    <A::Operator as Magma>::Set: Clone,
463    <A::Operand as Magma>::Set: Clone,
464{
465    fn drop(&mut self) {
466        let Self { index, tree, elt } = self;
467        let i = *index;
468        let elt = std::mem::replace(elt, tree.action.operand().id());
469        tree.buf.borrow_mut()[tree.len + i] = elt;
470        tree.build_range(i, i + 1);
471    }
472}
473
474impl<A> Deref for GetMutIndex<'_, A>
475where
476    A: MonoidAction,
477    <A::Operator as Magma>::Set: Clone,
478    <A::Operand as Magma>::Set: Clone,
479{
480    type Target = <A::Operand as Magma>::Set;
481    fn deref(&self) -> &Self::Target { &self.elt }
482}
483
484impl<A> DerefMut for GetMutIndex<'_, A>
485where
486    A: MonoidAction,
487    <A::Operator as Magma>::Set: Clone,
488    <A::Operand as Magma>::Set: Clone,
489{
490    fn deref_mut(&mut self) -> &mut Self::Target { &mut self.elt }
491}
492
493#[test]
494fn test_ancestors() {
495    use op_closure::OpClosure;
496    use op_closure_on_op_closure::OpClosureOnOpClosure;
497
498    for n in 1..=100 {
499        let mut height = vec![0; n + n];
500        for i in n..n + n {
501            height[i] = 1;
502        }
503        for i in (1..n).rev() {
504            if height[i << 1] == height[i << 1 | 1] {
505                height[i] = height[i << 1] + 1;
506            }
507        }
508
509        let action = OpClosureOnOpClosure::new(
510            OpClosure::new(|_, _| (), || ()),
511            OpClosure::new(|_, _| (), || ()),
512            |_, _| (),
513        );
514        let st: VecActSegtree<_> = (vec![(); n], action).into();
515        for l in 0..n {
516            for r in l..n {
517                let arch = {
518                    let mut arch = vec![];
519                    let mut l = n + l;
520                    let mut r = n + r + 1;
521                    while l < r {
522                        if l & 1 == 1 {
523                            arch.push(l);
524                            l += 1;
525                        }
526                        if r & 1 == 1 {
527                            r -= 1;
528                            arch.push(r);
529                        }
530                        l >>= 1;
531                        r >>= 1;
532                    }
533                    arch
534                };
535
536                let mut expected = vec![];
537                for mut i in arch {
538                    i >>= 1;
539                    while height[i] != 0 {
540                        expected.push(i);
541                        i >>= 1;
542                    }
543                }
544                expected.sort_unstable();
545                expected.dedup();
546
547                let actual = st.ancestors_sorted(l, r + 1);
548                assert_eq!(actual, expected);
549            }
550        }
551    }
552}
553
554#[cfg(test)]
555#[allow(dead_code)]
556fn dump(n: usize, node: &[usize]) {
557    use std::collections::BTreeSet;
558    let node: BTreeSet<_> = node.iter().copied().collect();
559
560    let len = n.next_power_of_two() << 1;
561    for i in 1..n + n {
562        let shift = WORD_SIZE - 1 - i.leading_zeros();
563        let ch = if node.contains(&i) { "o" } else { "-" };
564        eprint!("[{}]", ch.repeat((len >> shift) - 2));
565        if i + 1 == n + n || (i + 1).is_power_of_two() {
566            eprintln!();
567        }
568    }
569}
570
571#[test]
572fn test_fold_bisect() {
573    for n in (0..=32).chain(250..=260) {
574        test_fold_bisect_n(n);
575    }
576}
577
578#[cfg(test)]
579fn test_fold_bisect_n(n: usize) {
580    use op_affine_on_op_add_count::OpAffineOnOpAddCount;
581
582    let init = (1, 1);
583    let a = vec![init; n];
584    let mut st: VecActSegtree<OpAffineOnOpAddCount<i128>> = a.into();
585
586    let range_n: Vec<_> =
587        range(n).into_iter().filter_map(std::convert::identity).collect();
588
589    let actions: Vec<_> = (range_n.iter().rev())
590        .map(|&(l, r)| {
591            let range = l..r;
592            let x1 = (r - l).trailing_zeros() + 1;
593            let x0 = l;
594            (range, (x1 as i128, x0 as i128))
595        })
596        .collect();
597
598    let mut naive = vec![init; n];
599    for &(ref range, elt) in &actions {
600        for i in range.clone() {
601            let x = naive[i].0;
602            let (a, b) = elt;
603            naive[i].0 = a * x + b;
604        }
605        st.act(range.clone(), elt);
606    }
607
608    assert_eq!(st.unforced().len(), range_n.len() - n);
609    assert_eq!(Vec::<_>::from(st.clone()), naive);
610
611    let acc = {
612        let mut acc = vec![0; n + 1];
613        for i in 0..n {
614            acc[i + 1] = acc[i] + naive[i].0;
615        }
616        acc
617    };
618
619    for l in 0..=n {
620        for r in l..=n {
621            let st = st.clone();
622
623            let pred = |&(x, _): &(i128, _)| x <= acc[r] - acc[l];
624            let (index, folded) = st.fold_bisect(l, pred);
625
626            // return value
627            assert_eq!(
628                (index, folded),
629                (r, (acc[r] - acc[l], (r - l) as i128))
630            );
631
632            // postcondition
633            assert!(pred(&st.fold(l..index)));
634            assert!(index == n || !pred(&st.fold(l..index + 1)));
635
636            // actual values
637            assert_eq!(Vec::<_>::from(st), naive);
638        }
639    }
640
641    for r in 0..=n {
642        for l in 0..=r {
643            // clone to save unforced-ness
644            let st = st.clone();
645
646            let pred = |&(x, _): &(i128, _)| x <= acc[r] - acc[l];
647            let (index, folded) = st.fold_bisect_rev(r, pred);
648
649            // return value
650            assert_eq!(
651                (index, folded),
652                (l, (acc[r] - acc[l], (r - l) as i128))
653            );
654
655            // postcondition
656            assert!(pred(&st.fold(index..r)));
657            assert!(index == 0 || !pred(&st.fold(index - 1..r)));
658
659            // actual values
660            assert_eq!(Vec::<_>::from(st), naive);
661        }
662    }
663}
664
665#[cfg(test)]
666fn range(n: usize) -> Vec<Option<(usize, usize)>> {
667    let mut a = vec![None; n + n];
668    for i in 0..n {
669        a[n + i] = Some((i, i + 1));
670    }
671    for i in (1..n).rev() {
672        a[i] = match (a[2 * i], a[2 * i + 1]) {
673            (Some((ll, lr)), Some((rl, rr))) if lr == rl => Some((ll, rr)),
674            _ => None,
675        }
676    }
677    a
678}
679
680#[test]
681fn test_get_mut() {
682    for n in (0..=32).chain(1020..=1030) {
683        test_get_mut_n(n);
684    }
685}
686
687#[cfg(test)]
688fn test_get_mut_n(n: usize) {
689    use op_affine_on_op_add_count::OpAffineOnOpAddCount;
690
691    let init = (1, 1);
692    let a = vec![init; n];
693    let mut st: VecActSegtree<OpAffineOnOpAddCount<i128>> = a.into();
694
695    let range_n: Vec<_> =
696        range(n).into_iter().filter_map(std::convert::identity).collect();
697
698    let actions: Vec<_> = (range_n.iter().rev())
699        .map(|&(l, r)| {
700            let range = l..r;
701            let x1 = (r - l).trailing_zeros() + 1;
702            let x0 = l;
703            (range, (x1 as i128, x0 as i128))
704        })
705        .collect();
706
707    let mut naive = vec![init; n];
708    for &(ref range, elt) in &actions {
709        for i in range.clone() {
710            let x = naive[i].0;
711            let (a, b) = elt;
712            naive[i].0 = a * x + b;
713        }
714        st.act(range.clone(), elt);
715    }
716
717    for i in 0..n {
718        // clone to save unforced-ness
719        let mut st = st.clone();
720        let mut naive = naive.clone();
721
722        st.get_mut(i).unwrap().0 = 0;
723        naive[i].0 = 0;
724
725        assert_eq!(Vec::<_>::from(st), naive);
726    }
727}
nekolib/ds/vec_act_segtree.rs

nekolib/ds/
vec_act_segtree.rs
