competitive/math/
relaxed_convolution.rs

1use super::{ConvolveSteps, Zero};
2use std::{
3    fmt::Debug,
4    iter::zip,
5    marker::PhantomData,
6    ops::{AddAssign, Index},
7    slice::SliceIndex,
8};
9
10pub struct RelaxedConvolution<T, C>
11where
12    C: ConvolveSteps<T = Vec<T>>,
13{
14    a: Vec<T>,
15    b: Vec<T>,
16    c: Vec<T>,
17    _marker: PhantomData<fn() -> C>,
18}
19
20impl<T: Debug, C> Debug for RelaxedConvolution<T, C>
21where
22    C: ConvolveSteps<T = Vec<T>>,
23{
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("RelaxedConvolution")
26            .field("a", &self.a)
27            .field("b", &self.b)
28            .field("c", &self.c)
29            .finish()
30    }
31}
32
33impl<T, C> Default for RelaxedConvolution<T, C>
34where
35    C: ConvolveSteps<T = Vec<T>>,
36{
37    fn default() -> Self {
38        Self {
39            a: Default::default(),
40            b: Default::default(),
41            c: Default::default(),
42            _marker: PhantomData,
43        }
44    }
45}
46
47impl<T, C> RelaxedConvolution<T, C>
48where
49    T: Clone + Zero + AddAssign<T>,
50    C: ConvolveSteps<T = Vec<T>>,
51{
52    pub fn push(&mut self, x: T, y: T) {
53        let q = self.a.len();
54        self.a.push(x);
55        self.b.push(y);
56        self.c.push(T::zero());
57        if q != 0 {
58            self.c.push(T::zero());
59        }
60        let k = (q + 2).trailing_zeros();
61        let mut s = 0;
62        for k in 0..k + 1 - (1 << k == q + 2) as u32 {
63            let size = 1 << k;
64            self.calc_block(s, q + 1 - size, size);
65            if q + 1 - size != s {
66                self.calc_block(q + 1 - size, s, size);
67            }
68            s += size;
69        }
70    }
71
72    fn calc_block(&mut self, la: usize, lb: usize, size: usize) {
73        let a = self.a[la..la + size].to_vec();
74        let b = self.b[lb..lb + size].to_vec();
75        let c = C::convolve(a, b);
76        for (c, d) in zip(c, &mut self.c[la + lb..]) {
77            *d += c;
78        }
79    }
80}
81
82impl<T, C, I> Index<I> for RelaxedConvolution<T, C>
83where
84    C: ConvolveSteps<T = Vec<T>>,
85    I: SliceIndex<[T]>,
86{
87    type Output = <I as SliceIndex<[T]>>::Output;
88    fn index(&self, index: I) -> &Self::Output {
89        self.c.index(index)
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::{math::Convolve998244353, num::montgomery::MInt998244353, tools::Xorshift};
97
98    #[test]
99    fn test_relaxed_convolution() {
100        let mut rng = Xorshift::default();
101        for _ in 0..100 {
102            let n = rng.random(1..100);
103            let a: Vec<MInt998244353> = (0..n).map(|_| rng.random(..)).take(n).collect();
104            let b: Vec<MInt998244353> = (0..n).map(|_| rng.random(..)).take(n).collect();
105            let mut conv = RelaxedConvolution::<MInt998244353, Convolve998244353>::default();
106            for (&x, &y) in zip(&a, &b) {
107                conv.push(x, y);
108            }
109            let c = Convolve998244353::convolve(a, b);
110            for i in 0..n {
111                assert_eq!(conv[i], c[i]);
112            }
113        }
114    }
115}
competitive/math/relaxed_convolution.rs

competitive/math/
relaxed_convolution.rs
