Skip to main content

competitive/math/formal_power_series/
mod.rs

1use super::{
2    Convolve998244353, ConvolveSteps, MInt, MIntConvert, MIntConvolve, MemorizedFactorial,
3    NttReuse, One, PartialIgnoredOrd, Zero, berlekamp_massey, montgomery::MInt998244353,
4};
5use std::{
6    fmt::{self, Debug},
7    marker::PhantomData,
8    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
9};
10
11#[derive(Default)]
12pub struct FormalPowerSeries<T, C> {
13    pub data: Vec<T>,
14    _marker: PhantomData<C>,
15}
16
17impl<T, C> Debug for FormalPowerSeries<T, C>
18where
19    T: Debug,
20{
21    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
22        Debug::fmt(&self.data, f)
23    }
24}
25
26pub type Fps998244353 = FormalPowerSeries<MInt998244353, Convolve998244353>;
27pub type Fps<M> = FormalPowerSeries<MInt<M>, MIntConvolve<M>>;
28
29pub trait FormalPowerSeriesCoefficient:
30    Sized
31    + Clone
32    + Zero
33    + PartialEq
34    + One
35    + From<usize>
36    + From<isize>
37    + Add<Output = Self>
38    + Sub<Output = Self>
39    + Mul<Output = Self>
40    + Div<Output = Self>
41    + for<'r> Add<&'r Self, Output = Self>
42    + for<'r> Sub<&'r Self, Output = Self>
43    + for<'r> Mul<&'r Self, Output = Self>
44    + for<'r> Div<&'r Self, Output = Self>
45    + AddAssign<Self>
46    + SubAssign<Self>
47    + MulAssign<Self>
48    + DivAssign<Self>
49    + for<'r> AddAssign<&'r Self>
50    + for<'r> SubAssign<&'r Self>
51    + for<'r> MulAssign<&'r Self>
52    + for<'r> DivAssign<&'r Self>
53    + Neg<Output = Self>
54{
55    type Base: MIntConvert<usize> + MIntConvert<isize>;
56    fn pow(self, exp: usize) -> Self;
57    fn signed_pow(self, exp: isize) -> Self {
58        if exp >= 0 {
59            self.pow(exp as usize)
60        } else {
61            Self::one() / self.pow((-exp) as usize)
62        }
63    }
64    fn memorized_factorial(n: usize) -> MemorizedFactorial<Self::Base> {
65        MemorizedFactorial::new(n)
66    }
67    fn memorized_fact(mf: &MemorizedFactorial<Self::Base>) -> &[Self];
68    fn memorized_inv_fact(mf: &MemorizedFactorial<Self::Base>) -> &[Self];
69    fn memorized_inv(mf: &MemorizedFactorial<Self::Base>, n: usize) -> Self;
70}
71
72impl<M> FormalPowerSeriesCoefficient for MInt<M>
73where
74    M: MIntConvert<usize> + MIntConvert<isize>,
75{
76    type Base = M;
77    fn pow(self, exp: usize) -> Self {
78        Self::pow(self, exp)
79    }
80    fn memorized_fact(mf: &MemorizedFactorial<Self::Base>) -> &[Self] {
81        &mf.fact
82    }
83    fn memorized_inv_fact(mf: &MemorizedFactorial<Self::Base>) -> &[Self] {
84        &mf.inv_fact
85    }
86    fn memorized_inv(mf: &MemorizedFactorial<Self::Base>, n: usize) -> Self {
87        mf.inv(n)
88    }
89}
90
91pub trait FormalPowerSeriesCoefficientSqrt: FormalPowerSeriesCoefficient {
92    fn sqrt_coefficient(&self) -> Option<Self>;
93}
94
95impl<M> FormalPowerSeriesCoefficientSqrt for MInt<M>
96where
97    M: MIntConvert<u32> + MIntConvert<usize> + MIntConvert<isize>,
98{
99    fn sqrt_coefficient(&self) -> Option<Self> {
100        self.sqrt()
101    }
102}
103
104mod formal_power_series_impls;
105mod formal_power_series_nums;