fn add<T, C>(
left: FormalPowerSeries<T, C>,
right: FormalPowerSeries<T, C>,
) -> FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,Examples found in repository?
crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 58)
52 fn multiply_vector(
53 &self,
54 p: &FormalPowerSeries<T, C>,
55 q: &FormalPowerSeries<T, C>,
56 ) -> (FormalPowerSeries<T, C>, FormalPowerSeries<T, C>) {
57 (
58 add(&self.a00 * p, &self.a01 * q),
59 add(&self.a10 * p, &self.a11 * q),
60 )
61 }
62
63 fn left_multiply_step(&mut self, quotient: &[T]) {
64 swap(&mut self.a00, &mut self.a10);
65 swap(&mut self.a01, &mut self.a11);
66 let quotient: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(quotient.to_vec());
67 self.a10 = add(
68 replace(&mut self.a10, FormalPowerSeries::zero()),
69 quotient.clone() * &self.a00,
70 );
71 self.a11 = add(
72 replace(&mut self.a11, FormalPowerSeries::zero()),
73 quotient * &self.a01,
74 );
75 }
76
77 fn transform(&self, length: usize) -> FrequencyMatrix<C> {
78 FrequencyMatrix {
79 a00: reduced_transform(&self.a00, length),
80 a01: reduced_transform(&self.a01, length),
81 a10: reduced_transform(&self.a10, length),
82 a11: reduced_transform(&self.a11, length),
83 }
84 }
85
86 fn extend_transform(&self, frequency: FrequencyMatrix<C>, length: usize) -> FrequencyMatrix<C> {
87 fn extend<T, C>(fps: &FormalPowerSeries<T, C>, frequency: C::F, length: usize) -> C::F
88 where
89 T: FormalPowerSeriesCoefficient,
90 C: NttReuse<T = Vec<T>>,
91 C::F: Clone,
92 {
93 if fps.length() <= length / 2 {
94 C::ntt_doubling(frequency, false)
95 } else {
96 reduced_transform(fps, length)
97 }
98 }
99
100 FrequencyMatrix {
101 a00: extend(&self.a00, frequency.a00, length),
102 a01: extend(&self.a01, frequency.a01, length),
103 a10: extend(&self.a10, frequency.a10, length),
104 a11: extend(&self.a11, frequency.a11, length),
105 }
106 }
107}
108
109impl<T, C> FrequencyMatrix<C>
110where
111 T: FormalPowerSeriesCoefficient,
112 C: NttReuse<T = Vec<T>>,
113 C::F: Clone,
114{
115 fn product_sum(left_a: &C::F, right_a: &C::F, left_b: &C::F, right_b: &C::F) -> C::F {
116 let mut result = left_a.clone();
117 C::multiply_prefix(&mut result, right_a);
118 C::multiply_add(&mut result, left_b, right_b);
119 result
120 }
121
122 fn multiply(&self, right: &Self) -> Self {
123 Self {
124 a00: Self::product_sum(&self.a00, &right.a00, &self.a01, &right.a10),
125 a01: Self::product_sum(&self.a00, &right.a01, &self.a01, &right.a11),
126 a10: Self::product_sum(&self.a10, &right.a00, &self.a11, &right.a10),
127 a11: Self::product_sum(&self.a10, &right.a01, &self.a11, &right.a11),
128 }
129 }
130
131 fn apply(&self, p: &C::F, q: &C::F, length: usize) -> (Vec<T>, Vec<T>) {
132 (
133 C::inverse_transform_ntt(Self::product_sum(p, &self.a00, q, &self.a01), length),
134 C::inverse_transform_ntt(Self::product_sum(p, &self.a10, q, &self.a11), length),
135 )
136 }
137
138 fn left_multiply_step(self, quotient: &FormalPowerSeries<T, C>, length: usize) -> Self {
139 let negative_quotient = reduced_transform(&(-quotient), length);
140 let mut a10 = self.a00;
141 C::multiply_add(&mut a10, &negative_quotient, &self.a10);
142 let mut a11 = self.a01;
143 C::multiply_add(&mut a11, &negative_quotient, &self.a11);
144 let result = Self {
145 a00: self.a10,
146 a01: self.a11,
147 a10,
148 a11,
149 };
150 if C::MULTIPLE {
151 result.inverse_transform(length).transform(length)
152 } else {
153 result
154 }
155 }
156
157 fn inverse_transform(self, length: usize) -> FpsMatrix<T, C> {
158 FpsMatrix {
159 a00: FormalPowerSeries::from_vec(C::inverse_transform_ntt(self.a00, length)),
160 a01: FormalPowerSeries::from_vec(C::inverse_transform_ntt(self.a01, length)),
161 a10: FormalPowerSeries::from_vec(C::inverse_transform_ntt(self.a10, length)),
162 a11: FormalPowerSeries::from_vec(C::inverse_transform_ntt(self.a11, length)),
163 }
164 }
165}
166
167fn berlekamp_massey_naive<T>(a: &[T], max_work: usize) -> Option<Vec<T>>
168where
169 T: FormalPowerSeriesCoefficient,
170{
171 let n = a.len();
172 let mut b = Vec::with_capacity(n + 1);
173 let mut c = Vec::with_capacity(n + 1);
174 let mut temporary = Vec::with_capacity(n + 1);
175 b.push(T::one());
176 c.push(T::one());
177 let mut y = T::one();
178 let mut work = 0usize;
179 for k in 1..=n {
180 let c_len = c.len();
181 work = work.saturating_add(c_len);
182 if work > max_work {
183 return None;
184 }
185 let mut x = T::zero();
186 for (c, a) in c.iter().zip(&a[k - c_len..]) {
187 x += c.clone() * a.clone();
188 }
189 b.push(T::zero());
190 let b_len = b.len();
191 if x.is_zero() {
192 continue;
193 }
194 let frequency = x.clone() / y.clone();
195 if c_len < b_len {
196 swap(&mut c, &mut temporary);
197 c.clear();
198 c.resize_with(b_len - c_len, T::zero);
199 c.extend(temporary.iter().cloned());
200 for (c, b) in c.iter_mut().rev().zip(b.iter().rev()) {
201 *c -= frequency.clone() * b.clone();
202 }
203 swap(&mut b, &mut temporary);
204 y = x;
205 } else {
206 for (c, b) in c.iter_mut().rev().zip(b.iter().rev()) {
207 *c -= frequency.clone() * b.clone();
208 }
209 }
210 }
211 c.reverse();
212 Some(c)
213}
214
215impl<T, C> FormalPowerSeries<T, C>
216where
217 T: FormalPowerSeriesCoefficient,
218 C: NttReuse<T = Vec<T>>,
219 C::F: Clone,
220{
221 pub fn berlekamp_massey(input: &[T]) -> Self {
222 if input.last().is_none_or(|value| value.is_zero())
223 && input.iter().all(|value| value.is_zero())
224 {
225 return Self::one();
226 }
227 let max_work = if input.len() <= 1536 {
228 usize::MAX
229 } else {
230 input.len().saturating_mul(2)
231 };
232 if let Some(recurrence) = berlekamp_massey_naive(input, max_work) {
233 return Self::from_vec(recurrence);
234 }
235 let n = input.len();
236 let leading_zeros = input.iter().take_while(|value| value.is_zero()).count();
237 let sequence = Self::from_vec(input.to_vec()).trimed();
238 let mut modulus = Self::zeros(n + 1);
239 modulus[n] = T::one();
240 let (matrix, _) = half_gcd(&modulus, &sequence, n / 2, n.max(1).next_power_of_two());
241 let (x, y) = matrix.multiply_vector(&modulus, &sequence);
242 let mut recurrence = if y.length() == 0 {
243 matrix.a01.clone()
244 } else {
245 matrix.a11.clone()
246 };
247 let recurrence_leading_zeros = recurrence
248 .iter()
249 .take_while(|value| value.is_zero())
250 .count();
251 if recurrence_leading_zeros > 0 {
252 let (division, _) = x.div_rem(y.clone());
253 recurrence = add(recurrence * division, matrix.a01);
254 }
255 let inverse = T::one() / &recurrence[0];
256 for value in recurrence.iter_mut() {
257 *value *= &inverse;
258 }
259 let minimum_length = (leading_zeros + 2).max(y.length() + 1);
260 if recurrence.length() < minimum_length {
261 recurrence.resize(minimum_length);
262 }
263 recurrence
264 }