haar_lib/ds/
foldable_deque.rs

1//! 半群で畳み込み可能なdeque
2pub use crate::algebra::traits::*;
3
4#[derive(Clone, Default, Debug)]
5/// 半群で畳み込み可能なdeque
6pub struct FoldableDeque<S: Semigroup> {
7    semigroup: S,
8    front_stack: Vec<S::Element>,
9    back_stack: Vec<S::Element>,
10    front_sum: Vec<S::Element>,
11    back_sum: Vec<S::Element>,
12}
13
14impl<S: Semigroup> FoldableDeque<S>
15where
16    S::Element: Clone,
17{
18    /// 空の`FoldableDeque<S>`を生成する。
19    pub fn new(semigroup: S) -> Self {
20        Self {
21            semigroup,
22            front_stack: vec![],
23            back_stack: vec![],
24            front_sum: vec![],
25            back_sum: vec![],
26        }
27    }
28
29    fn f(&self, a: Option<S::Element>, b: Option<S::Element>) -> Option<S::Element> {
30        match (a, b) {
31            (Some(a), Some(b)) => Some(self.semigroup.op(a, b)),
32            (x, None) => x,
33            (None, x) => x,
34        }
35    }
36
37    /// すべての要素を`S`の演算で畳み込んだ結果を返す。
38    pub fn fold(&self) -> Option<S::Element> {
39        self.f(
40            self.front_sum.last().cloned(),
41            self.back_sum.last().cloned(),
42        )
43    }
44
45    /// 末尾に`value`を追加する。
46    pub fn push_back(&mut self, value: S::Element) {
47        self.back_stack.push(value.clone());
48        self.back_sum
49            .push(self.f(self.back_sum.last().cloned(), Some(value)).unwrap());
50    }
51
52    /// 先頭に`value`を追加する。
53    pub fn push_front(&mut self, value: S::Element) {
54        self.front_stack.push(value.clone());
55        self.front_sum
56            .push(self.f(Some(value), self.front_sum.last().cloned()).unwrap());
57    }
58
59    fn build_sum(&mut self) {
60        for value in &self.front_stack {
61            self.front_sum.push(
62                self.f(Some(value.clone()), self.front_sum.last().cloned())
63                    .unwrap(),
64            );
65        }
66
67        for value in &self.back_stack {
68            self.back_sum.push(
69                self.f(self.back_sum.last().cloned(), Some(value.clone()))
70                    .unwrap(),
71            );
72        }
73    }
74
75    /// 先頭の要素を削除して返す。
76    pub fn pop_front(&mut self) -> Option<S::Element> {
77        if self.front_stack.is_empty() {
78            self.back_sum.clear();
79
80            let n = self.back_stack.len();
81            if n == 0 {
82                return None;
83            }
84
85            self.front_stack = self.back_stack.split_off(n.div_ceil(2));
86            std::mem::swap(&mut self.front_stack, &mut self.back_stack);
87            self.front_stack.reverse();
88
89            self.build_sum();
90        }
91
92        self.front_sum.pop();
93        self.front_stack.pop()
94    }
95
96    /// 末尾の要素を削除して返す。
97    pub fn pop_back(&mut self) -> Option<S::Element> {
98        if self.back_stack.is_empty() {
99            self.front_sum.clear();
100
101            let n = self.front_stack.len();
102            if n == 0 {
103                return None;
104            }
105
106            self.back_stack = self.front_stack.split_off(n.div_ceil(2));
107            std::mem::swap(&mut self.front_stack, &mut self.back_stack);
108            self.back_stack.reverse();
109
110            self.build_sum();
111        }
112
113        self.back_sum.pop();
114        self.back_stack.pop()
115    }
116
117    /// 先頭の要素への参照を返す。
118    pub fn front(&self) -> Option<&S::Element> {
119        self.front_stack.last().or_else(|| self.back_stack.first())
120    }
121
122    /// 末尾の要素への参照を返す。
123    pub fn back(&self) -> Option<&S::Element> {
124        self.back_stack.last().or_else(|| self.front_stack.first())
125    }
126
127    /// 要素数を返す。
128    pub fn len(&self) -> usize {
129        self.front_stack.len() + self.back_stack.len()
130    }
131
132    /// 要素数が`0`なら`true`を返す。
133    pub fn is_empty(&self) -> bool {
134        self.front_stack.is_empty() && self.back_stack.is_empty()
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use crate::{
142        algebra::{affine::*, semiring::add_mul_mod::AddMulMod},
143        math::{linear::Linear, prime_mod::Prime},
144        num::const_modint::{ConstModInt, ConstModIntBuilder},
145    };
146    use rand::Rng;
147    use std::collections::VecDeque;
148
149    const M: u32 = 998244353;
150    type Mint = ConstModInt<Prime<M>>;
151
152    #[test]
153    fn test() {
154        let mut rng = rand::thread_rng();
155
156        let mut deq = VecDeque::<Linear<Mint>>::new();
157
158        let m = Composition(AddMulMod(ConstModIntBuilder::new()));
159        let mut swag = FoldableDeque::new(m);
160
161        for _ in 0..1000 {
162            assert_eq!(deq.front(), swag.front());
163            assert_eq!(deq.back(), swag.back());
164            assert_eq!(deq.len(), swag.len());
165
166            let ty = rng.gen_range(0..5);
167
168            match ty {
169                0 => {
170                    let a = Mint::new(rng.gen_range(0..M));
171                    let b = Mint::new(rng.gen_range(0..M));
172                    deq.push_front(Linear::new(a, b));
173                    swag.push_front(Linear::new(a, b));
174                }
175                1 => {
176                    let a = Mint::new(rng.gen_range(0..M));
177                    let b = Mint::new(rng.gen_range(0..M));
178                    deq.push_back(Linear::new(a, b));
179                    swag.push_back(Linear::new(a, b));
180                }
181                2 => {
182                    assert_eq!(deq.pop_front(), swag.pop_front());
183                }
184                3 => {
185                    assert_eq!(deq.pop_back(), swag.pop_back());
186                }
187                4 => {
188                    assert_eq!(
189                        deq.iter().cloned().fold_m(&m),
190                        swag.fold().unwrap_or(m.id())
191                    );
192                }
193                _ => unreachable!(),
194            }
195        }
196    }
197}