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

haar_lib/ds/
foldable_deque.rs
