pub struct FormalPowerSeries<T, C> {
pub data: Vec<T>,
_marker: PhantomData<C>,
}Fields§
§data: Vec<T>§_marker: PhantomData<C>Implementations§
Source§impl<T, C> FormalPowerSeries<T, C>
impl<T, C> FormalPowerSeries<T, C>
Sourcepub fn berlekamp_massey(input: &[T]) -> Self
pub fn berlekamp_massey(input: &[T]) -> Self
Examples found in repository?
More examples
crates/competitive/src/math/black_box_mint_matrix.rs (line 27)
14 fn minimal_polynomial(&self) -> Vec<MInt<M>> {
15 assert_eq!(self.shape().0, self.shape().1);
16 let n = self.shape().0;
17 let mut rng = Xorshift::new();
18 let b: Vec<MInt<M>> = (0..n).map(|_| MInt::from(rng.rand64())).collect();
19 let u: Vec<MInt<M>> = (0..n).map(|_| MInt::from(rng.rand64())).collect();
20 let a: Vec<MInt<M>> = (0..2 * n)
21 .scan(b, |b, _| {
22 let a = MInt::dot_product(b, &u);
23 *b = self.apply(b);
24 Some(a)
25 })
26 .collect();
27 let polynomial: Fps<M> = FormalPowerSeries::berlekamp_massey(&a);
28 let mut p = polynomial.data;
29 p.reverse();
30 p
31 }Source§impl<T, C> FormalPowerSeries<T, C>
impl<T, C> FormalPowerSeries<T, C>
Sourcepub fn from_vec(data: Vec<T>) -> Self
pub fn from_vec(data: Vec<T>) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 38)
37 fn clone(&self) -> Self {
38 Self::from_vec(self.data.clone())
39 }
40}
41impl<T, C> PartialEq for FormalPowerSeries<T, C>
42where
43 T: PartialEq,
44{
45 fn eq(&self, other: &Self) -> bool {
46 self.data.eq(&other.data)
47 }
48}
49impl<T, C> Eq for FormalPowerSeries<T, C> where T: PartialEq {}
50
51impl<T, C> FormalPowerSeries<T, C>
52where
53 T: Zero,
54{
55 pub fn zeros(deg: usize) -> Self {
56 repeat_with(T::zero).take(deg).collect()
57 }
58 pub fn resize(&mut self, deg: usize) {
59 self.data.resize_with(deg, Zero::zero)
60 }
61 pub fn resized(mut self, deg: usize) -> Self {
62 self.resize(deg);
63 self
64 }
65 pub fn reversed(mut self) -> Self {
66 self.data.reverse();
67 self
68 }
69}
70
71impl<T, C> FormalPowerSeries<T, C>
72where
73 T: Zero + Clone,
74{
75 pub fn coeff(&self, deg: usize) -> T {
76 self.data.get(deg).cloned().unwrap_or_else(T::zero)
77 }
78}
79
80impl<T, C> FormalPowerSeries<T, C>
81where
82 T: Zero + PartialEq,
83{
84 pub fn trim_tail_zeros(&mut self) {
85 let mut len = self.length();
86 while len > 0 {
87 if self.data[len - 1].is_zero() {
88 len -= 1;
89 } else {
90 break;
91 }
92 }
93 self.truncate(len);
94 }
95 pub fn trimed(mut self) -> Self {
96 self.trim_tail_zeros();
97 self
98 }
99}
100
101impl<T, C> Zero for FormalPowerSeries<T, C>
102where
103 T: PartialEq,
104{
105 fn zero() -> Self {
106 Self::from_vec(Vec::new())
107 }
108}
109impl<T, C> One for FormalPowerSeries<T, C>
110where
111 T: PartialEq + One,
112{
113 fn one() -> Self {
114 Self::from(T::one())
115 }
116}
117
118impl<T, C> IntoIterator for FormalPowerSeries<T, C> {
119 type Item = T;
120 type IntoIter = std::vec::IntoIter<T>;
121 fn into_iter(self) -> Self::IntoIter {
122 self.data.into_iter()
123 }
124}
125impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C> {
126 type Item = &'a T;
127 type IntoIter = Iter<'a, T>;
128 fn into_iter(self) -> Self::IntoIter {
129 self.data.iter()
130 }
131}
132impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C> {
133 type Item = &'a mut T;
134 type IntoIter = IterMut<'a, T>;
135 fn into_iter(self) -> Self::IntoIter {
136 self.data.iter_mut()
137 }
138}
139
140impl<T, C> FromIterator<T> for FormalPowerSeries<T, C> {
141 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
142 Self::from_vec(iter.into_iter().collect())
143 }
144}
145
146impl<T, C> Index<usize> for FormalPowerSeries<T, C> {
147 type Output = T;
148 fn index(&self, index: usize) -> &Self::Output {
149 &self.data[index]
150 }
151}
152impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C> {
153 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
154 &mut self.data[index]
155 }
156}
157
158impl<T, C> From<T> for FormalPowerSeries<T, C> {
159 fn from(x: T) -> Self {
160 once(x).collect()
161 }
162}
163impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C> {
164 fn from(data: Vec<T>) -> Self {
165 Self::from_vec(data)
166 }
167}
168
169impl<T, C> FormalPowerSeries<T, C>
170where
171 T: FormalPowerSeriesCoefficient,
172{
173 pub fn prefix_ref(&self, deg: usize) -> Self {
174 if deg < self.length() {
175 Self::from_vec(self.data[..deg].to_vec())
176 } else {
177 self.clone()
178 }
179 }
180 pub fn prefix(mut self, deg: usize) -> Self {
181 self.data.truncate(deg);
182 self
183 }
184 pub fn even(mut self) -> Self {
185 let mut keep = false;
186 self.data.retain(|_| {
187 keep = !keep;
188 keep
189 });
190 self
191 }
192 pub fn odd(mut self) -> Self {
193 let mut keep = true;
194 self.data.retain(|_| {
195 keep = !keep;
196 keep
197 });
198 self
199 }
200 pub fn diff(mut self) -> Self {
201 let mut c = T::one();
202 for i in 1..self.length() {
203 self.data[i - 1] = self.data[i].clone() * &c;
204 c += T::one();
205 }
206 self.data.pop();
207 self
208 }
209 pub fn integral(mut self) -> Self {
210 let n = self.length();
211 let mut fact = Vec::with_capacity(n + 1);
212 let mut c = T::one();
213 fact.push(c.clone());
214 for _ in 1..n {
215 fact.push(fact.last().cloned().unwrap() * c.clone());
216 c += T::one();
217 }
218 let mut invf = T::one() / (fact.last().cloned().unwrap() * c.clone());
219 self.data.push(T::zero());
220 for i in (1..=n).rev() {
221 self.data[i] = self.data[i - 1].clone() * (invf.clone() * fact.pop().unwrap());
222 invf *= c.clone();
223 c -= T::one();
224 }
225 self.data[0] = T::zero();
226 self
227 }
228 pub fn parity_inversion(mut self) -> Self {
229 self.iter_mut()
230 .skip(1)
231 .step_by(2)
232 .for_each(|x| *x = -x.clone());
233 self
234 }
235 pub fn eval(&self, x: T) -> T {
236 self.iter()
237 .rev()
238 .fold(T::zero(), |sum, a| x.clone() * sum + a.clone())
239 }
240}
241
242impl<T, C> FormalPowerSeries<T, C>
243where
244 T: FormalPowerSeriesCoefficient,
245 C: ConvolveSteps<T = Vec<T>>,
246{
247 #[inline]
248 fn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize> {
249 let limit = deg.next_power_of_two().trailing_zeros() as usize * factor;
250 let mut count = 0;
251 let mut step = 0;
252 for (i, value) in self.iter().take(deg).enumerate() {
253 if value.is_zero() {
254 continue;
255 }
256 count += 1;
257 if step != 1 {
258 step = gcd(step, i as u64);
259 }
260 if count > limit {
261 return None;
262 }
263 }
264 Some(step.max(1) as usize)
265 }
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }
1367 /// f(x) <- f(x + a)
1368 pub fn taylor_shift(mut self, a: T) -> Self {
1369 let f = T::memorized_factorial(self.length());
1370 let n = self.length();
1371 for (i, coef) in self.data.iter_mut().enumerate() {
1372 *coef *= T::memorized_fact(&f)[i].clone();
1373 }
1374 self.data.reverse();
1375 let mut b = a.clone();
1376 let mut g = Self::from_vec(T::memorized_inv_fact(&f)[..n].to_vec());
1377 for i in 1..n {
1378 g[i] *= b.clone();
1379 b *= a.clone();
1380 }
1381 self *= g;
1382 self.truncate(n);
1383 self.data.reverse();
1384 for (i, coef) in self.data.iter_mut().enumerate() {
1385 *coef *= T::memorized_inv_fact(&f)[i].clone();
1386 }
1387 self
1388 }More examples
Additional examples can be found in:
- crates/library_checker/src/polynomial/polynomial_taylor_shift.rs
- crates/library_checker/src/polynomial/multipoint_evaluation.rs
- crates/library_checker/src/polynomial/compositional_inverse_of_formal_power_series.rs
- crates/library_checker/src/polynomial/compositional_inverse_of_formal_power_series_large.rs
- crates/library_checker/src/polynomial/sqrt_of_formal_power_series.rs
- crates/library_checker/src/other/kth_term_of_linearly_recurrent_sequence.rs
- crates/library_checker/src/polynomial/division_of_polynomials.rs
- crates/library_checker/src/polynomial/exp_of_formal_power_series_sparse.rs
- crates/library_checker/src/polynomial/inv_of_formal_power_series_sparse.rs
- crates/library_checker/src/polynomial/log_of_formal_power_series_sparse.rs
- crates/library_checker/src/polynomial/pow_of_formal_power_series_sparse.rs
- crates/library_checker/src/enumerative_combinatorics/sharp_p_subset_sum.rs
- crates/library_checker/src/polynomial/sqrt_of_formal_power_series_sparse.rs
- crates/competitive/src/math/formal_power_series/berlekamp_massey.rs
- crates/competitive/src/math/black_box_mint_matrix.rs
Sourcepub fn length(&self) -> usize
pub fn length(&self) -> usize
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_nums.rs (line 15)
14 fn add_assign(&mut self, rhs: T) {
15 if self.length() == 0 {
16 self.data.push(T::zero());
17 }
18 self.data[0].add_assign(rhs);
19 }
20}
21impl<T, C> SubAssign<T> for FormalPowerSeries<T, C>
22where
23 T: FormalPowerSeriesCoefficient,
24{
25 fn sub_assign(&mut self, rhs: T) {
26 if self.length() == 0 {
27 self.data.push(T::zero());
28 }
29 self.data[0].sub_assign(rhs);
30 self.trim_tail_zeros();
31 }
32}
33impl<T, C> MulAssign<T> for FormalPowerSeries<T, C>
34where
35 T: FormalPowerSeriesCoefficient,
36{
37 fn mul_assign(&mut self, rhs: T) {
38 for x in self.iter_mut() {
39 x.mul_assign(&rhs);
40 }
41 }
42}
43impl<T, C> DivAssign<T> for FormalPowerSeries<T, C>
44where
45 T: FormalPowerSeriesCoefficient,
46{
47 fn div_assign(&mut self, rhs: T) {
48 let rinv = T::one() / rhs;
49 for x in self.iter_mut() {
50 x.mul_assign(&rinv);
51 }
52 }
53}
54macro_rules! impl_fps_single_binop {
55 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
56 impl<T, C> $imp_assign<&T> for FormalPowerSeries<T, C>
57 where
58 T: FormalPowerSeriesCoefficient,
59 {
60 fn $method_assign(&mut self, rhs: &T) {
61 $imp_assign::$method_assign(self, rhs.clone());
62 }
63 }
64 impl<T, C> $imp<T> for FormalPowerSeries<T, C>
65 where
66 T: FormalPowerSeriesCoefficient,
67 {
68 type Output = Self;
69 fn $method(mut self, rhs: T) -> Self::Output {
70 $imp_assign::$method_assign(&mut self, rhs);
71 self
72 }
73 }
74 impl<T, C> $imp<&T> for FormalPowerSeries<T, C>
75 where
76 T: FormalPowerSeriesCoefficient,
77 {
78 type Output = Self;
79 fn $method(mut self, rhs: &T) -> Self::Output {
80 $imp_assign::$method_assign(&mut self, rhs);
81 self
82 }
83 }
84 impl<T, C> $imp<T> for &FormalPowerSeries<T, C>
85 where
86 T: FormalPowerSeriesCoefficient,
87 {
88 type Output = FormalPowerSeries<T, C>;
89 fn $method(self, rhs: T) -> Self::Output {
90 $imp::$method(self.clone(), rhs)
91 }
92 }
93 impl<T, C> $imp<&T> for &FormalPowerSeries<T, C>
94 where
95 T: FormalPowerSeriesCoefficient,
96 {
97 type Output = FormalPowerSeries<T, C>;
98 fn $method(self, rhs: &T) -> Self::Output {
99 $imp::$method(self.clone(), rhs)
100 }
101 }
102 };
103}
104impl_fps_single_binop!(Add, add, AddAssign, add_assign);
105impl_fps_single_binop!(Sub, sub, SubAssign, sub_assign);
106impl_fps_single_binop!(Mul, mul, MulAssign, mul_assign);
107impl_fps_single_binop!(Div, div, DivAssign, div_assign);
108
109impl<T, C> AddAssign<&Self> for FormalPowerSeries<T, C>
110where
111 T: FormalPowerSeriesCoefficient,
112{
113 fn add_assign(&mut self, rhs: &Self) {
114 if self.length() < rhs.length() {
115 self.resize(rhs.length());
116 }
117 for (x, y) in self.iter_mut().zip(rhs.iter()) {
118 x.add_assign(y);
119 }
120 }
121}
122impl<T, C> SubAssign<&Self> for FormalPowerSeries<T, C>
123where
124 T: FormalPowerSeriesCoefficient,
125{
126 fn sub_assign(&mut self, rhs: &Self) {
127 if self.length() < rhs.length() {
128 self.resize(rhs.length());
129 }
130 for (x, y) in self.iter_mut().zip(rhs.iter()) {
131 x.sub_assign(y);
132 }
133 self.trim_tail_zeros();
134 }
135}
136
137macro_rules! impl_fps_binop_addsub {
138 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
139 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
140 where
141 T: FormalPowerSeriesCoefficient,
142 {
143 fn $method_assign(&mut self, rhs: Self) {
144 $imp_assign::$method_assign(self, &rhs);
145 }
146 }
147 impl<T, C> $imp for FormalPowerSeries<T, C>
148 where
149 T: FormalPowerSeriesCoefficient,
150 {
151 type Output = Self;
152 fn $method(mut self, rhs: Self) -> Self::Output {
153 $imp_assign::$method_assign(&mut self, &rhs);
154 self
155 }
156 }
157 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
158 where
159 T: FormalPowerSeriesCoefficient,
160 {
161 type Output = Self;
162 fn $method(mut self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
163 $imp_assign::$method_assign(&mut self, rhs);
164 self
165 }
166 }
167 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
168 where
169 T: FormalPowerSeriesCoefficient,
170 {
171 type Output = FormalPowerSeries<T, C>;
172 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
173 let mut self_ = self.clone();
174 $imp_assign::$method_assign(&mut self_, &rhs);
175 self_
176 }
177 }
178 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
179 where
180 T: FormalPowerSeriesCoefficient,
181 {
182 type Output = FormalPowerSeries<T, C>;
183 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
184 let mut self_ = self.clone();
185 $imp_assign::$method_assign(&mut self_, rhs);
186 self_
187 }
188 }
189 };
190}
191impl_fps_binop_addsub!(Add, add, AddAssign, add_assign);
192impl_fps_binop_addsub!(Sub, sub, SubAssign, sub_assign);
193
194impl<T, C> Mul for FormalPowerSeries<T, C>
195where
196 C: ConvolveSteps<T = Vec<T>>,
197{
198 type Output = Self;
199 fn mul(self, rhs: Self) -> Self::Output {
200 Self::from_vec(C::convolve(self.data, rhs.data))
201 }
202}
203impl<T, C> Div for FormalPowerSeries<T, C>
204where
205 T: FormalPowerSeriesCoefficient,
206 C: ConvolveSteps<T = Vec<T>>,
207{
208 type Output = Self;
209 fn div(mut self, mut rhs: Self) -> Self::Output {
210 self.trim_tail_zeros();
211 rhs.trim_tail_zeros();
212 if self.length() < rhs.length() {
213 return Self::zero();
214 }
215 self.data.reverse();
216 rhs.data.reverse();
217 let n = self.length() - rhs.length() + 1;
218 let mut res = self * rhs.inv(n);
219 res.truncate(n);
220 res.data.reverse();
221 res
222 }
223}
224impl<T, C> Rem for FormalPowerSeries<T, C>
225where
226 T: FormalPowerSeriesCoefficient,
227 C: ConvolveSteps<T = Vec<T>>,
228{
229 type Output = Self;
230 fn rem(self, rhs: Self) -> Self::Output {
231 let mut rem = self.clone() - self / rhs.clone() * rhs;
232 rem.trim_tail_zeros();
233 rem
234 }
235}
236
237impl<T, C> FormalPowerSeries<T, C>
238where
239 T: FormalPowerSeriesCoefficient,
240 C: ConvolveSteps<T = Vec<T>>,
241{
242 pub fn div_rem(self, rhs: Self) -> (Self, Self) {
243 let div = self.clone() / rhs.clone();
244 let mut rem = self - div.clone() * rhs;
245 rem.trim_tail_zeros();
246 (div, rem)
247 }
248}
249
250macro_rules! impl_fps_binop_conv {
251 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
252 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
253 where
254 T: FormalPowerSeriesCoefficient,
255 C: ConvolveSteps<T = Vec<T>>,
256 {
257 fn $method_assign(&mut self, rhs: Self) {
258 *self = $imp::$method(Self::from_vec(take(&mut self.data)), rhs);
259 }
260 }
261 impl<T, C> $imp_assign<&Self> for FormalPowerSeries<T, C>
262 where
263 T: FormalPowerSeriesCoefficient,
264 C: ConvolveSteps<T = Vec<T>>,
265 {
266 fn $method_assign(&mut self, rhs: &Self) {
267 $imp_assign::$method_assign(self, rhs.clone());
268 }
269 }
270 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
271 where
272 T: FormalPowerSeriesCoefficient,
273 C: ConvolveSteps<T = Vec<T>>,
274 {
275 type Output = Self;
276 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
277 $imp::$method(self, rhs.clone())
278 }
279 }
280 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
281 where
282 T: FormalPowerSeriesCoefficient,
283 C: ConvolveSteps<T = Vec<T>>,
284 {
285 type Output = FormalPowerSeries<T, C>;
286 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
287 $imp::$method(self.clone(), rhs)
288 }
289 }
290 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
291 where
292 T: FormalPowerSeriesCoefficient,
293 C: ConvolveSteps<T = Vec<T>>,
294 {
295 type Output = FormalPowerSeries<T, C>;
296 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
297 $imp::$method(self.clone(), rhs.clone())
298 }
299 }
300 };
301}
302impl_fps_binop_conv!(Mul, mul, MulAssign, mul_assign);
303impl_fps_binop_conv!(Div, div, DivAssign, div_assign);
304impl_fps_binop_conv!(Rem, rem, RemAssign, rem_assign);
305
306impl<T, C> Neg for FormalPowerSeries<T, C>
307where
308 T: FormalPowerSeriesCoefficient,
309{
310 type Output = Self;
311 fn neg(mut self) -> Self::Output {
312 for x in self.iter_mut() {
313 *x = -x.clone();
314 }
315 self
316 }
317}
318impl<T, C> Neg for &FormalPowerSeries<T, C>
319where
320 T: FormalPowerSeriesCoefficient,
321{
322 type Output = FormalPowerSeries<T, C>;
323 fn neg(self) -> Self::Output {
324 self.clone().neg()
325 }
326}
327
328impl<T, C> ShrAssign<usize> for FormalPowerSeries<T, C>
329where
330 T: FormalPowerSeriesCoefficient,
331{
332 fn shr_assign(&mut self, rhs: usize) {
333 if self.length() <= rhs {
334 *self = Self::zero();
335 } else {
336 for i in rhs..self.length() {
337 self[i - rhs] = self[i].clone();
338 }
339 self.truncate(self.length() - rhs);
340 }
341 }
342}
343impl<T, C> ShlAssign<usize> for FormalPowerSeries<T, C>
344where
345 T: FormalPowerSeriesCoefficient,
346{
347 fn shl_assign(&mut self, rhs: usize) {
348 let n = self.length();
349 self.resize(n + rhs);
350 for i in (0..n).rev() {
351 self[i + rhs] = self[i].clone();
352 }
353 for i in 0..rhs {
354 self[i] = T::zero();
355 }
356 }
357}
358
359impl<T, C> Shr<usize> for FormalPowerSeries<T, C>
360where
361 T: FormalPowerSeriesCoefficient,
362{
363 type Output = Self;
364 fn shr(mut self, rhs: usize) -> Self::Output {
365 self.shr_assign(rhs);
366 self
367 }
368}
369impl<T, C> Shl<usize> for FormalPowerSeries<T, C>
370where
371 T: FormalPowerSeriesCoefficient,
372{
373 type Output = Self;
374 fn shl(mut self, rhs: usize) -> Self::Output {
375 self.shl_assign(rhs);
376 self
377 }
378}
379impl<T, C> Shr<usize> for &FormalPowerSeries<T, C>
380where
381 T: FormalPowerSeriesCoefficient,
382{
383 type Output = FormalPowerSeries<T, C>;
384 fn shr(self, rhs: usize) -> Self::Output {
385 if self.length() <= rhs {
386 Self::Output::zero()
387 } else {
388 let mut f = Self::Output::zeros(self.length() - rhs);
389 for i in rhs..self.length() {
390 f[i - rhs] = self[i].clone();
391 }
392 f
393 }
394 }
395}
396impl<T, C> Shl<usize> for &FormalPowerSeries<T, C>
397where
398 T: FormalPowerSeriesCoefficient,
399{
400 type Output = FormalPowerSeries<T, C>;
401 fn shl(self, rhs: usize) -> Self::Output {
402 let mut f = Self::Output::zeros(self.length() + rhs);
403 for (i, x) in self.iter().cloned().enumerate().rev() {
404 f[i + rhs] = x;
405 }
406 f
407 }More examples
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 85)
84 pub fn trim_tail_zeros(&mut self) {
85 let mut len = self.length();
86 while len > 0 {
87 if self.data[len - 1].is_zero() {
88 len -= 1;
89 } else {
90 break;
91 }
92 }
93 self.truncate(len);
94 }
95 pub fn trimed(mut self) -> Self {
96 self.trim_tail_zeros();
97 self
98 }
99}
100
101impl<T, C> Zero for FormalPowerSeries<T, C>
102where
103 T: PartialEq,
104{
105 fn zero() -> Self {
106 Self::from_vec(Vec::new())
107 }
108}
109impl<T, C> One for FormalPowerSeries<T, C>
110where
111 T: PartialEq + One,
112{
113 fn one() -> Self {
114 Self::from(T::one())
115 }
116}
117
118impl<T, C> IntoIterator for FormalPowerSeries<T, C> {
119 type Item = T;
120 type IntoIter = std::vec::IntoIter<T>;
121 fn into_iter(self) -> Self::IntoIter {
122 self.data.into_iter()
123 }
124}
125impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C> {
126 type Item = &'a T;
127 type IntoIter = Iter<'a, T>;
128 fn into_iter(self) -> Self::IntoIter {
129 self.data.iter()
130 }
131}
132impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C> {
133 type Item = &'a mut T;
134 type IntoIter = IterMut<'a, T>;
135 fn into_iter(self) -> Self::IntoIter {
136 self.data.iter_mut()
137 }
138}
139
140impl<T, C> FromIterator<T> for FormalPowerSeries<T, C> {
141 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
142 Self::from_vec(iter.into_iter().collect())
143 }
144}
145
146impl<T, C> Index<usize> for FormalPowerSeries<T, C> {
147 type Output = T;
148 fn index(&self, index: usize) -> &Self::Output {
149 &self.data[index]
150 }
151}
152impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C> {
153 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
154 &mut self.data[index]
155 }
156}
157
158impl<T, C> From<T> for FormalPowerSeries<T, C> {
159 fn from(x: T) -> Self {
160 once(x).collect()
161 }
162}
163impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C> {
164 fn from(data: Vec<T>) -> Self {
165 Self::from_vec(data)
166 }
167}
168
169impl<T, C> FormalPowerSeries<T, C>
170where
171 T: FormalPowerSeriesCoefficient,
172{
173 pub fn prefix_ref(&self, deg: usize) -> Self {
174 if deg < self.length() {
175 Self::from_vec(self.data[..deg].to_vec())
176 } else {
177 self.clone()
178 }
179 }
180 pub fn prefix(mut self, deg: usize) -> Self {
181 self.data.truncate(deg);
182 self
183 }
184 pub fn even(mut self) -> Self {
185 let mut keep = false;
186 self.data.retain(|_| {
187 keep = !keep;
188 keep
189 });
190 self
191 }
192 pub fn odd(mut self) -> Self {
193 let mut keep = true;
194 self.data.retain(|_| {
195 keep = !keep;
196 keep
197 });
198 self
199 }
200 pub fn diff(mut self) -> Self {
201 let mut c = T::one();
202 for i in 1..self.length() {
203 self.data[i - 1] = self.data[i].clone() * &c;
204 c += T::one();
205 }
206 self.data.pop();
207 self
208 }
209 pub fn integral(mut self) -> Self {
210 let n = self.length();
211 let mut fact = Vec::with_capacity(n + 1);
212 let mut c = T::one();
213 fact.push(c.clone());
214 for _ in 1..n {
215 fact.push(fact.last().cloned().unwrap() * c.clone());
216 c += T::one();
217 }
218 let mut invf = T::one() / (fact.last().cloned().unwrap() * c.clone());
219 self.data.push(T::zero());
220 for i in (1..=n).rev() {
221 self.data[i] = self.data[i - 1].clone() * (invf.clone() * fact.pop().unwrap());
222 invf *= c.clone();
223 c -= T::one();
224 }
225 self.data[0] = T::zero();
226 self
227 }
228 pub fn parity_inversion(mut self) -> Self {
229 self.iter_mut()
230 .skip(1)
231 .step_by(2)
232 .for_each(|x| *x = -x.clone());
233 self
234 }
235 pub fn eval(&self, x: T) -> T {
236 self.iter()
237 .rev()
238 .fold(T::zero(), |sum, a| x.clone() * sum + a.clone())
239 }
240}
241
242impl<T, C> FormalPowerSeries<T, C>
243where
244 T: FormalPowerSeriesCoefficient,
245 C: ConvolveSteps<T = Vec<T>>,
246{
247 #[inline]
248 fn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize> {
249 let limit = deg.next_power_of_two().trailing_zeros() as usize * factor;
250 let mut count = 0;
251 let mut step = 0;
252 for (i, value) in self.iter().take(deg).enumerate() {
253 if value.is_zero() {
254 continue;
255 }
256 count += 1;
257 if step != 1 {
258 step = gcd(step, i as u64);
259 }
260 if count > limit {
261 return None;
262 }
263 }
264 Some(step.max(1) as usize)
265 }
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }
1367 /// f(x) <- f(x + a)
1368 pub fn taylor_shift(mut self, a: T) -> Self {
1369 let f = T::memorized_factorial(self.length());
1370 let n = self.length();
1371 for (i, coef) in self.data.iter_mut().enumerate() {
1372 *coef *= T::memorized_fact(&f)[i].clone();
1373 }
1374 self.data.reverse();
1375 let mut b = a.clone();
1376 let mut g = Self::from_vec(T::memorized_inv_fact(&f)[..n].to_vec());
1377 for i in 1..n {
1378 g[i] *= b.clone();
1379 b *= a.clone();
1380 }
1381 self *= g;
1382 self.truncate(n);
1383 self.data.reverse();
1384 for (i, coef) in self.data.iter_mut().enumerate() {
1385 *coef *= T::memorized_inv_fact(&f)[i].clone();
1386 }
1387 self
1388 }crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 93)
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 }
265}
266
267fn degree<T, C>(fps: &FormalPowerSeries<T, C>) -> isize {
268 fps.length() as isize - 1
269}
270
271fn add<T, C>(
272 left: FormalPowerSeries<T, C>,
273 right: FormalPowerSeries<T, C>,
274) -> FormalPowerSeries<T, C>
275where
276 T: FormalPowerSeriesCoefficient,
277{
278 (left + right).trimed()
279}
280
281fn tail<T, C>(fps: &FormalPowerSeries<T, C>, start: isize) -> FormalPowerSeries<T, C>
282where
283 T: FormalPowerSeriesCoefficient,
284{
285 let start = start.max(0) as usize;
286 if start >= fps.length() {
287 FormalPowerSeries::zero()
288 } else {
289 FormalPowerSeries::from_vec(fps.data[start..].to_vec())
290 }
291}
292
293fn coefficient<T, C>(fps: &FormalPowerSeries<T, C>, index: isize) -> T
294where
295 T: FormalPowerSeriesCoefficient,
296{
297 if index < 0 {
298 T::zero()
299 } else {
300 fps.coeff(index as usize)
301 }
302}
303
304fn brute_force<T, C>(
305 mut p: FormalPowerSeries<T, C>,
306 mut q: FormalPowerSeries<T, C>,
307 k: usize,
308) -> FpsMatrix<T, C>
309where
310 T: FormalPowerSeriesCoefficient,
311 C: NttReuse<T = Vec<T>>,
312 C::F: Clone,
313{
314 let threshold = degree(&p) - k as isize;
315 let mut matrix = FpsMatrix::identity();
316 while q.length() as isize > threshold {
317 let q_degree = q.length() - 1;
318 let mut negative_quotient = vec![T::zero(); p.length() - q.length() + 1];
319 let inverse = -T::one() / &q[q_degree];
320 for i in (0..negative_quotient.len()).rev() {
321 negative_quotient[i] = p[i + q_degree].clone() * &inverse;
322 p[i + q_degree] = T::zero();
323 for j in 0..q_degree {
324 let value = negative_quotient[i].clone() * &q[j];
325 p[i + j] += &value;
326 }
327 }
328 matrix.left_multiply_step(&negative_quotient);
329 p.truncate(q_degree);
330 p.trim_tail_zeros();
331 swap(&mut p, &mut q);
332 }
333 matrix
334}
335
336fn reduced_transform<T, C>(fps: &FormalPowerSeries<T, C>, length: usize) -> C::F
337where
338 T: FormalPowerSeriesCoefficient,
339 C: NttReuse<T = Vec<T>>,
340{
341 let mut coefficients = vec![T::zero(); length];
342 for (i, value) in fps.iter().enumerate() {
343 coefficients[i & (length - 1)] += value;
344 }
345 C::transform_ntt(coefficients, length)
346}
347
348fn transform_window<T, C>(fps: &FormalPowerSeries<T, C>, end: isize, length: usize) -> C::F
349where
350 T: FormalPowerSeriesCoefficient,
351 C: NttReuse<T = Vec<T>>,
352{
353 let start = end - length as isize;
354 let coefficients = (start..end).map(|index| coefficient(fps, index)).collect();
355 C::transform_ntt(coefficients, length)
356}
357
358fn half_gcd<T, C>(
359 p: &FormalPowerSeries<T, C>,
360 q: &FormalPowerSeries<T, C>,
361 k: usize,
362 length: usize,
363) -> (FpsMatrix<T, C>, FrequencyMatrix<C>)
364where
365 T: FormalPowerSeriesCoefficient,
366 C: NttReuse<T = Vec<T>>,
367 C::F: Clone,
368{
369 let d = degree(p);
370 if degree(q) < d - k as isize {
371 let matrix = FpsMatrix::identity();
372 let frequency = matrix.transform(length);
373 return (matrix, frequency);
374 }
375 if k == 1 {
376 let matrix = FpsMatrix {
377 a00: FormalPowerSeries::zero(),
378 a01: FormalPowerSeries::one(),
379 a10: FormalPowerSeries::one(),
380 a11: -(tail(p, d - 2) / tail(q, d - 2)),
381 };
382 let frequency = matrix.transform(length);
383 return (matrix, frequency);
384 }
385 if p.length().min(q.length()) <= 32 {
386 let matrix = brute_force(p.clone(), q.clone(), k);
387 let frequency = matrix.transform(length);
388 return (matrix, frequency);
389 }
390
391 let half = length / 2;
392 if k <= half {
393 let (matrix, frequency) = half_gcd(p, q, k, half);
394 let frequency = matrix.extend_transform(frequency, length);
395 return (matrix, frequency);
396 }
397
398 let (matrix, mut matrix_frequency) = half_gcd(
399 &tail(p, d - 2 * half as isize),
400 &tail(q, d - 2 * half as isize),
401 half,
402 length,
403 );
404 let degeneracy = half as isize - degree(&matrix.a11);
405
406 let (p0, q0) = matrix_frequency.apply(
407 &transform_window(p, d - half as isize + degeneracy, length),
408 &transform_window(q, d - half as isize + degeneracy, length),
409 length,
410 );
411 let (p1, q1) = matrix_frequency.apply(
412 &transform_window(p, d - 2 * half as isize, length),
413 &transform_window(q, d - 2 * half as isize, length),
414 length,
415 );
416 let part_length = (half as isize + degeneracy) as usize;
417 let mut p_reduced = p1[length - part_length..].to_vec();
418 p_reduced.extend_from_slice(&p0[length - part_length..]);
419 let mut q_reduced = q1[length - part_length..].to_vec();
420 q_reduced.extend_from_slice(&q0[length - part_length..]);
421 let mut q_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(q_reduced).trimed();
422
423 let position = d - half as isize + degeneracy;
424 let mut leading = T::zero();
425 for i in 0..=position {
426 leading += coefficient(p, i) * coefficient(&matrix.a00, position - i)
427 + coefficient(q, i) * coefficient(&matrix.a01, position - i);
428 }
429 p_reduced.push(leading);
430 let mut p_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(p_reduced);
431 if degree(&q_reduced) < 3 * half as isize + degeneracy - k as isize {
432 return (matrix, matrix_frequency);
433 }
434
435 let mut remaining = k as isize - degree(&matrix.a11);
436 let mut top_product = matrix.a11.data.last().unwrap().clone();
437 let mut product_degree = degree(&matrix.a11);
438 if degeneracy > 0 {
439 let skip = (2 * half as isize + 2 * degeneracy - (d - half as isize + degeneracy)).max(0);
440 let (division, remainder) = tail(&p_reduced, skip).div_rem(tail(&q_reduced, skip));
441 remaining -= degree(&division);
442 top_product *= -division.data.last().unwrap().clone();
443 product_degree += degree(&division);
444 matrix_frequency = matrix_frequency.left_multiply_step(&division, length);
445 swap(&mut p_reduced, &mut q_reduced);
446 q_reduced = FormalPowerSeries::zeros(skip as usize);
447 q_reduced.data.extend(remainder.data);
448 }
449
450 let start = 3 * half as isize + degeneracy - k as isize - remaining;
451 let (right_matrix, right_frequency) = half_gcd(
452 &tail(&p_reduced, start),
453 &tail(&q_reduced, start),
454 remaining as usize,
455 length,
456 );
457 let product_frequency = right_frequency.multiply(&matrix_frequency);
458 let mut product = product_frequency.clone().inverse_transform(length);
459 product.a00.truncate(k);
460 product.a00.trim_tail_zeros();
461 product.a01.truncate(k);
462 product.a01.trim_tail_zeros();
463 product.a10.truncate(k);
464 product.a10.trim_tail_zeros();
465 product_degree += degree(&right_matrix.a11);
466 if product_degree == length as isize {
467 product.a11.resize(k + 1);
468 let highest = top_product * right_matrix.a11.data.last().unwrap();
469 product.a11[k] = highest.clone();
470 product.a11[0] -= highest;
471 }
472 product.a11.trim_tail_zeros();
473 let product_frequency = if C::MULTIPLE {
474 product.transform(length)
475 } else {
476 product_frequency
477 };
478 (product, product_frequency)
479}Sourcepub fn truncate(&mut self, deg: usize)
pub fn truncate(&mut self, deg: usize)
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 93)
84 pub fn trim_tail_zeros(&mut self) {
85 let mut len = self.length();
86 while len > 0 {
87 if self.data[len - 1].is_zero() {
88 len -= 1;
89 } else {
90 break;
91 }
92 }
93 self.truncate(len);
94 }
95 pub fn trimed(mut self) -> Self {
96 self.trim_tail_zeros();
97 self
98 }
99}
100
101impl<T, C> Zero for FormalPowerSeries<T, C>
102where
103 T: PartialEq,
104{
105 fn zero() -> Self {
106 Self::from_vec(Vec::new())
107 }
108}
109impl<T, C> One for FormalPowerSeries<T, C>
110where
111 T: PartialEq + One,
112{
113 fn one() -> Self {
114 Self::from(T::one())
115 }
116}
117
118impl<T, C> IntoIterator for FormalPowerSeries<T, C> {
119 type Item = T;
120 type IntoIter = std::vec::IntoIter<T>;
121 fn into_iter(self) -> Self::IntoIter {
122 self.data.into_iter()
123 }
124}
125impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C> {
126 type Item = &'a T;
127 type IntoIter = Iter<'a, T>;
128 fn into_iter(self) -> Self::IntoIter {
129 self.data.iter()
130 }
131}
132impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C> {
133 type Item = &'a mut T;
134 type IntoIter = IterMut<'a, T>;
135 fn into_iter(self) -> Self::IntoIter {
136 self.data.iter_mut()
137 }
138}
139
140impl<T, C> FromIterator<T> for FormalPowerSeries<T, C> {
141 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
142 Self::from_vec(iter.into_iter().collect())
143 }
144}
145
146impl<T, C> Index<usize> for FormalPowerSeries<T, C> {
147 type Output = T;
148 fn index(&self, index: usize) -> &Self::Output {
149 &self.data[index]
150 }
151}
152impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C> {
153 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
154 &mut self.data[index]
155 }
156}
157
158impl<T, C> From<T> for FormalPowerSeries<T, C> {
159 fn from(x: T) -> Self {
160 once(x).collect()
161 }
162}
163impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C> {
164 fn from(data: Vec<T>) -> Self {
165 Self::from_vec(data)
166 }
167}
168
169impl<T, C> FormalPowerSeries<T, C>
170where
171 T: FormalPowerSeriesCoefficient,
172{
173 pub fn prefix_ref(&self, deg: usize) -> Self {
174 if deg < self.length() {
175 Self::from_vec(self.data[..deg].to_vec())
176 } else {
177 self.clone()
178 }
179 }
180 pub fn prefix(mut self, deg: usize) -> Self {
181 self.data.truncate(deg);
182 self
183 }
184 pub fn even(mut self) -> Self {
185 let mut keep = false;
186 self.data.retain(|_| {
187 keep = !keep;
188 keep
189 });
190 self
191 }
192 pub fn odd(mut self) -> Self {
193 let mut keep = true;
194 self.data.retain(|_| {
195 keep = !keep;
196 keep
197 });
198 self
199 }
200 pub fn diff(mut self) -> Self {
201 let mut c = T::one();
202 for i in 1..self.length() {
203 self.data[i - 1] = self.data[i].clone() * &c;
204 c += T::one();
205 }
206 self.data.pop();
207 self
208 }
209 pub fn integral(mut self) -> Self {
210 let n = self.length();
211 let mut fact = Vec::with_capacity(n + 1);
212 let mut c = T::one();
213 fact.push(c.clone());
214 for _ in 1..n {
215 fact.push(fact.last().cloned().unwrap() * c.clone());
216 c += T::one();
217 }
218 let mut invf = T::one() / (fact.last().cloned().unwrap() * c.clone());
219 self.data.push(T::zero());
220 for i in (1..=n).rev() {
221 self.data[i] = self.data[i - 1].clone() * (invf.clone() * fact.pop().unwrap());
222 invf *= c.clone();
223 c -= T::one();
224 }
225 self.data[0] = T::zero();
226 self
227 }
228 pub fn parity_inversion(mut self) -> Self {
229 self.iter_mut()
230 .skip(1)
231 .step_by(2)
232 .for_each(|x| *x = -x.clone());
233 self
234 }
235 pub fn eval(&self, x: T) -> T {
236 self.iter()
237 .rev()
238 .fold(T::zero(), |sum, a| x.clone() * sum + a.clone())
239 }
240}
241
242impl<T, C> FormalPowerSeries<T, C>
243where
244 T: FormalPowerSeriesCoefficient,
245 C: ConvolveSteps<T = Vec<T>>,
246{
247 #[inline]
248 fn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize> {
249 let limit = deg.next_power_of_two().trailing_zeros() as usize * factor;
250 let mut count = 0;
251 let mut step = 0;
252 for (i, value) in self.iter().take(deg).enumerate() {
253 if value.is_zero() {
254 continue;
255 }
256 count += 1;
257 if step != 1 {
258 step = gcd(step, i as u64);
259 }
260 if count > limit {
261 return None;
262 }
263 }
264 Some(step.max(1) as usize)
265 }
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }
1367 /// f(x) <- f(x + a)
1368 pub fn taylor_shift(mut self, a: T) -> Self {
1369 let f = T::memorized_factorial(self.length());
1370 let n = self.length();
1371 for (i, coef) in self.data.iter_mut().enumerate() {
1372 *coef *= T::memorized_fact(&f)[i].clone();
1373 }
1374 self.data.reverse();
1375 let mut b = a.clone();
1376 let mut g = Self::from_vec(T::memorized_inv_fact(&f)[..n].to_vec());
1377 for i in 1..n {
1378 g[i] *= b.clone();
1379 b *= a.clone();
1380 }
1381 self *= g;
1382 self.truncate(n);
1383 self.data.reverse();
1384 for (i, coef) in self.data.iter_mut().enumerate() {
1385 *coef *= T::memorized_inv_fact(&f)[i].clone();
1386 }
1387 self
1388 }More examples
crates/competitive/src/math/formal_power_series/formal_power_series_nums.rs (line 219)
209 fn div(mut self, mut rhs: Self) -> Self::Output {
210 self.trim_tail_zeros();
211 rhs.trim_tail_zeros();
212 if self.length() < rhs.length() {
213 return Self::zero();
214 }
215 self.data.reverse();
216 rhs.data.reverse();
217 let n = self.length() - rhs.length() + 1;
218 let mut res = self * rhs.inv(n);
219 res.truncate(n);
220 res.data.reverse();
221 res
222 }
223}
224impl<T, C> Rem for FormalPowerSeries<T, C>
225where
226 T: FormalPowerSeriesCoefficient,
227 C: ConvolveSteps<T = Vec<T>>,
228{
229 type Output = Self;
230 fn rem(self, rhs: Self) -> Self::Output {
231 let mut rem = self.clone() - self / rhs.clone() * rhs;
232 rem.trim_tail_zeros();
233 rem
234 }
235}
236
237impl<T, C> FormalPowerSeries<T, C>
238where
239 T: FormalPowerSeriesCoefficient,
240 C: ConvolveSteps<T = Vec<T>>,
241{
242 pub fn div_rem(self, rhs: Self) -> (Self, Self) {
243 let div = self.clone() / rhs.clone();
244 let mut rem = self - div.clone() * rhs;
245 rem.trim_tail_zeros();
246 (div, rem)
247 }
248}
249
250macro_rules! impl_fps_binop_conv {
251 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
252 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
253 where
254 T: FormalPowerSeriesCoefficient,
255 C: ConvolveSteps<T = Vec<T>>,
256 {
257 fn $method_assign(&mut self, rhs: Self) {
258 *self = $imp::$method(Self::from_vec(take(&mut self.data)), rhs);
259 }
260 }
261 impl<T, C> $imp_assign<&Self> for FormalPowerSeries<T, C>
262 where
263 T: FormalPowerSeriesCoefficient,
264 C: ConvolveSteps<T = Vec<T>>,
265 {
266 fn $method_assign(&mut self, rhs: &Self) {
267 $imp_assign::$method_assign(self, rhs.clone());
268 }
269 }
270 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
271 where
272 T: FormalPowerSeriesCoefficient,
273 C: ConvolveSteps<T = Vec<T>>,
274 {
275 type Output = Self;
276 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
277 $imp::$method(self, rhs.clone())
278 }
279 }
280 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
281 where
282 T: FormalPowerSeriesCoefficient,
283 C: ConvolveSteps<T = Vec<T>>,
284 {
285 type Output = FormalPowerSeries<T, C>;
286 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
287 $imp::$method(self.clone(), rhs)
288 }
289 }
290 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
291 where
292 T: FormalPowerSeriesCoefficient,
293 C: ConvolveSteps<T = Vec<T>>,
294 {
295 type Output = FormalPowerSeries<T, C>;
296 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
297 $imp::$method(self.clone(), rhs.clone())
298 }
299 }
300 };
301}
302impl_fps_binop_conv!(Mul, mul, MulAssign, mul_assign);
303impl_fps_binop_conv!(Div, div, DivAssign, div_assign);
304impl_fps_binop_conv!(Rem, rem, RemAssign, rem_assign);
305
306impl<T, C> Neg for FormalPowerSeries<T, C>
307where
308 T: FormalPowerSeriesCoefficient,
309{
310 type Output = Self;
311 fn neg(mut self) -> Self::Output {
312 for x in self.iter_mut() {
313 *x = -x.clone();
314 }
315 self
316 }
317}
318impl<T, C> Neg for &FormalPowerSeries<T, C>
319where
320 T: FormalPowerSeriesCoefficient,
321{
322 type Output = FormalPowerSeries<T, C>;
323 fn neg(self) -> Self::Output {
324 self.clone().neg()
325 }
326}
327
328impl<T, C> ShrAssign<usize> for FormalPowerSeries<T, C>
329where
330 T: FormalPowerSeriesCoefficient,
331{
332 fn shr_assign(&mut self, rhs: usize) {
333 if self.length() <= rhs {
334 *self = Self::zero();
335 } else {
336 for i in rhs..self.length() {
337 self[i - rhs] = self[i].clone();
338 }
339 self.truncate(self.length() - rhs);
340 }
341 }crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 329)
304fn brute_force<T, C>(
305 mut p: FormalPowerSeries<T, C>,
306 mut q: FormalPowerSeries<T, C>,
307 k: usize,
308) -> FpsMatrix<T, C>
309where
310 T: FormalPowerSeriesCoefficient,
311 C: NttReuse<T = Vec<T>>,
312 C::F: Clone,
313{
314 let threshold = degree(&p) - k as isize;
315 let mut matrix = FpsMatrix::identity();
316 while q.length() as isize > threshold {
317 let q_degree = q.length() - 1;
318 let mut negative_quotient = vec![T::zero(); p.length() - q.length() + 1];
319 let inverse = -T::one() / &q[q_degree];
320 for i in (0..negative_quotient.len()).rev() {
321 negative_quotient[i] = p[i + q_degree].clone() * &inverse;
322 p[i + q_degree] = T::zero();
323 for j in 0..q_degree {
324 let value = negative_quotient[i].clone() * &q[j];
325 p[i + j] += &value;
326 }
327 }
328 matrix.left_multiply_step(&negative_quotient);
329 p.truncate(q_degree);
330 p.trim_tail_zeros();
331 swap(&mut p, &mut q);
332 }
333 matrix
334}
335
336fn reduced_transform<T, C>(fps: &FormalPowerSeries<T, C>, length: usize) -> C::F
337where
338 T: FormalPowerSeriesCoefficient,
339 C: NttReuse<T = Vec<T>>,
340{
341 let mut coefficients = vec![T::zero(); length];
342 for (i, value) in fps.iter().enumerate() {
343 coefficients[i & (length - 1)] += value;
344 }
345 C::transform_ntt(coefficients, length)
346}
347
348fn transform_window<T, C>(fps: &FormalPowerSeries<T, C>, end: isize, length: usize) -> C::F
349where
350 T: FormalPowerSeriesCoefficient,
351 C: NttReuse<T = Vec<T>>,
352{
353 let start = end - length as isize;
354 let coefficients = (start..end).map(|index| coefficient(fps, index)).collect();
355 C::transform_ntt(coefficients, length)
356}
357
358fn half_gcd<T, C>(
359 p: &FormalPowerSeries<T, C>,
360 q: &FormalPowerSeries<T, C>,
361 k: usize,
362 length: usize,
363) -> (FpsMatrix<T, C>, FrequencyMatrix<C>)
364where
365 T: FormalPowerSeriesCoefficient,
366 C: NttReuse<T = Vec<T>>,
367 C::F: Clone,
368{
369 let d = degree(p);
370 if degree(q) < d - k as isize {
371 let matrix = FpsMatrix::identity();
372 let frequency = matrix.transform(length);
373 return (matrix, frequency);
374 }
375 if k == 1 {
376 let matrix = FpsMatrix {
377 a00: FormalPowerSeries::zero(),
378 a01: FormalPowerSeries::one(),
379 a10: FormalPowerSeries::one(),
380 a11: -(tail(p, d - 2) / tail(q, d - 2)),
381 };
382 let frequency = matrix.transform(length);
383 return (matrix, frequency);
384 }
385 if p.length().min(q.length()) <= 32 {
386 let matrix = brute_force(p.clone(), q.clone(), k);
387 let frequency = matrix.transform(length);
388 return (matrix, frequency);
389 }
390
391 let half = length / 2;
392 if k <= half {
393 let (matrix, frequency) = half_gcd(p, q, k, half);
394 let frequency = matrix.extend_transform(frequency, length);
395 return (matrix, frequency);
396 }
397
398 let (matrix, mut matrix_frequency) = half_gcd(
399 &tail(p, d - 2 * half as isize),
400 &tail(q, d - 2 * half as isize),
401 half,
402 length,
403 );
404 let degeneracy = half as isize - degree(&matrix.a11);
405
406 let (p0, q0) = matrix_frequency.apply(
407 &transform_window(p, d - half as isize + degeneracy, length),
408 &transform_window(q, d - half as isize + degeneracy, length),
409 length,
410 );
411 let (p1, q1) = matrix_frequency.apply(
412 &transform_window(p, d - 2 * half as isize, length),
413 &transform_window(q, d - 2 * half as isize, length),
414 length,
415 );
416 let part_length = (half as isize + degeneracy) as usize;
417 let mut p_reduced = p1[length - part_length..].to_vec();
418 p_reduced.extend_from_slice(&p0[length - part_length..]);
419 let mut q_reduced = q1[length - part_length..].to_vec();
420 q_reduced.extend_from_slice(&q0[length - part_length..]);
421 let mut q_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(q_reduced).trimed();
422
423 let position = d - half as isize + degeneracy;
424 let mut leading = T::zero();
425 for i in 0..=position {
426 leading += coefficient(p, i) * coefficient(&matrix.a00, position - i)
427 + coefficient(q, i) * coefficient(&matrix.a01, position - i);
428 }
429 p_reduced.push(leading);
430 let mut p_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(p_reduced);
431 if degree(&q_reduced) < 3 * half as isize + degeneracy - k as isize {
432 return (matrix, matrix_frequency);
433 }
434
435 let mut remaining = k as isize - degree(&matrix.a11);
436 let mut top_product = matrix.a11.data.last().unwrap().clone();
437 let mut product_degree = degree(&matrix.a11);
438 if degeneracy > 0 {
439 let skip = (2 * half as isize + 2 * degeneracy - (d - half as isize + degeneracy)).max(0);
440 let (division, remainder) = tail(&p_reduced, skip).div_rem(tail(&q_reduced, skip));
441 remaining -= degree(&division);
442 top_product *= -division.data.last().unwrap().clone();
443 product_degree += degree(&division);
444 matrix_frequency = matrix_frequency.left_multiply_step(&division, length);
445 swap(&mut p_reduced, &mut q_reduced);
446 q_reduced = FormalPowerSeries::zeros(skip as usize);
447 q_reduced.data.extend(remainder.data);
448 }
449
450 let start = 3 * half as isize + degeneracy - k as isize - remaining;
451 let (right_matrix, right_frequency) = half_gcd(
452 &tail(&p_reduced, start),
453 &tail(&q_reduced, start),
454 remaining as usize,
455 length,
456 );
457 let product_frequency = right_frequency.multiply(&matrix_frequency);
458 let mut product = product_frequency.clone().inverse_transform(length);
459 product.a00.truncate(k);
460 product.a00.trim_tail_zeros();
461 product.a01.truncate(k);
462 product.a01.trim_tail_zeros();
463 product.a10.truncate(k);
464 product.a10.trim_tail_zeros();
465 product_degree += degree(&right_matrix.a11);
466 if product_degree == length as isize {
467 product.a11.resize(k + 1);
468 let highest = top_product * right_matrix.a11.data.last().unwrap();
469 product.a11[k] = highest.clone();
470 product.a11[0] -= highest;
471 }
472 product.a11.trim_tail_zeros();
473 let product_frequency = if C::MULTIPLE {
474 product.transform(length)
475 } else {
476 product_frequency
477 };
478 (product, product_frequency)
479}Sourcepub fn iter(&self) -> Iter<'_, T> ⓘ
pub fn iter(&self) -> Iter<'_, T> ⓘ
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 236)
235 pub fn eval(&self, x: T) -> T {
236 self.iter()
237 .rev()
238 .fold(T::zero(), |sum, a| x.clone() * sum + a.clone())
239 }
240}
241
242impl<T, C> FormalPowerSeries<T, C>
243where
244 T: FormalPowerSeriesCoefficient,
245 C: ConvolveSteps<T = Vec<T>>,
246{
247 #[inline]
248 fn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize> {
249 let limit = deg.next_power_of_two().trailing_zeros() as usize * factor;
250 let mut count = 0;
251 let mut step = 0;
252 for (i, value) in self.iter().take(deg).enumerate() {
253 if value.is_zero() {
254 continue;
255 }
256 count += 1;
257 if step != 1 {
258 step = gcd(step, i as u64);
259 }
260 if count > limit {
261 return None;
262 }
263 }
264 Some(step.max(1) as usize)
265 }
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }More examples
crates/competitive/src/math/formal_power_series/formal_power_series_nums.rs (line 117)
113 fn add_assign(&mut self, rhs: &Self) {
114 if self.length() < rhs.length() {
115 self.resize(rhs.length());
116 }
117 for (x, y) in self.iter_mut().zip(rhs.iter()) {
118 x.add_assign(y);
119 }
120 }
121}
122impl<T, C> SubAssign<&Self> for FormalPowerSeries<T, C>
123where
124 T: FormalPowerSeriesCoefficient,
125{
126 fn sub_assign(&mut self, rhs: &Self) {
127 if self.length() < rhs.length() {
128 self.resize(rhs.length());
129 }
130 for (x, y) in self.iter_mut().zip(rhs.iter()) {
131 x.sub_assign(y);
132 }
133 self.trim_tail_zeros();
134 }
135}
136
137macro_rules! impl_fps_binop_addsub {
138 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
139 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
140 where
141 T: FormalPowerSeriesCoefficient,
142 {
143 fn $method_assign(&mut self, rhs: Self) {
144 $imp_assign::$method_assign(self, &rhs);
145 }
146 }
147 impl<T, C> $imp for FormalPowerSeries<T, C>
148 where
149 T: FormalPowerSeriesCoefficient,
150 {
151 type Output = Self;
152 fn $method(mut self, rhs: Self) -> Self::Output {
153 $imp_assign::$method_assign(&mut self, &rhs);
154 self
155 }
156 }
157 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
158 where
159 T: FormalPowerSeriesCoefficient,
160 {
161 type Output = Self;
162 fn $method(mut self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
163 $imp_assign::$method_assign(&mut self, rhs);
164 self
165 }
166 }
167 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
168 where
169 T: FormalPowerSeriesCoefficient,
170 {
171 type Output = FormalPowerSeries<T, C>;
172 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
173 let mut self_ = self.clone();
174 $imp_assign::$method_assign(&mut self_, &rhs);
175 self_
176 }
177 }
178 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
179 where
180 T: FormalPowerSeriesCoefficient,
181 {
182 type Output = FormalPowerSeries<T, C>;
183 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
184 let mut self_ = self.clone();
185 $imp_assign::$method_assign(&mut self_, rhs);
186 self_
187 }
188 }
189 };
190}
191impl_fps_binop_addsub!(Add, add, AddAssign, add_assign);
192impl_fps_binop_addsub!(Sub, sub, SubAssign, sub_assign);
193
194impl<T, C> Mul for FormalPowerSeries<T, C>
195where
196 C: ConvolveSteps<T = Vec<T>>,
197{
198 type Output = Self;
199 fn mul(self, rhs: Self) -> Self::Output {
200 Self::from_vec(C::convolve(self.data, rhs.data))
201 }
202}
203impl<T, C> Div for FormalPowerSeries<T, C>
204where
205 T: FormalPowerSeriesCoefficient,
206 C: ConvolveSteps<T = Vec<T>>,
207{
208 type Output = Self;
209 fn div(mut self, mut rhs: Self) -> Self::Output {
210 self.trim_tail_zeros();
211 rhs.trim_tail_zeros();
212 if self.length() < rhs.length() {
213 return Self::zero();
214 }
215 self.data.reverse();
216 rhs.data.reverse();
217 let n = self.length() - rhs.length() + 1;
218 let mut res = self * rhs.inv(n);
219 res.truncate(n);
220 res.data.reverse();
221 res
222 }
223}
224impl<T, C> Rem for FormalPowerSeries<T, C>
225where
226 T: FormalPowerSeriesCoefficient,
227 C: ConvolveSteps<T = Vec<T>>,
228{
229 type Output = Self;
230 fn rem(self, rhs: Self) -> Self::Output {
231 let mut rem = self.clone() - self / rhs.clone() * rhs;
232 rem.trim_tail_zeros();
233 rem
234 }
235}
236
237impl<T, C> FormalPowerSeries<T, C>
238where
239 T: FormalPowerSeriesCoefficient,
240 C: ConvolveSteps<T = Vec<T>>,
241{
242 pub fn div_rem(self, rhs: Self) -> (Self, Self) {
243 let div = self.clone() / rhs.clone();
244 let mut rem = self - div.clone() * rhs;
245 rem.trim_tail_zeros();
246 (div, rem)
247 }
248}
249
250macro_rules! impl_fps_binop_conv {
251 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
252 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
253 where
254 T: FormalPowerSeriesCoefficient,
255 C: ConvolveSteps<T = Vec<T>>,
256 {
257 fn $method_assign(&mut self, rhs: Self) {
258 *self = $imp::$method(Self::from_vec(take(&mut self.data)), rhs);
259 }
260 }
261 impl<T, C> $imp_assign<&Self> for FormalPowerSeries<T, C>
262 where
263 T: FormalPowerSeriesCoefficient,
264 C: ConvolveSteps<T = Vec<T>>,
265 {
266 fn $method_assign(&mut self, rhs: &Self) {
267 $imp_assign::$method_assign(self, rhs.clone());
268 }
269 }
270 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
271 where
272 T: FormalPowerSeriesCoefficient,
273 C: ConvolveSteps<T = Vec<T>>,
274 {
275 type Output = Self;
276 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
277 $imp::$method(self, rhs.clone())
278 }
279 }
280 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
281 where
282 T: FormalPowerSeriesCoefficient,
283 C: ConvolveSteps<T = Vec<T>>,
284 {
285 type Output = FormalPowerSeries<T, C>;
286 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
287 $imp::$method(self.clone(), rhs)
288 }
289 }
290 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
291 where
292 T: FormalPowerSeriesCoefficient,
293 C: ConvolveSteps<T = Vec<T>>,
294 {
295 type Output = FormalPowerSeries<T, C>;
296 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
297 $imp::$method(self.clone(), rhs.clone())
298 }
299 }
300 };
301}
302impl_fps_binop_conv!(Mul, mul, MulAssign, mul_assign);
303impl_fps_binop_conv!(Div, div, DivAssign, div_assign);
304impl_fps_binop_conv!(Rem, rem, RemAssign, rem_assign);
305
306impl<T, C> Neg for FormalPowerSeries<T, C>
307where
308 T: FormalPowerSeriesCoefficient,
309{
310 type Output = Self;
311 fn neg(mut self) -> Self::Output {
312 for x in self.iter_mut() {
313 *x = -x.clone();
314 }
315 self
316 }
317}
318impl<T, C> Neg for &FormalPowerSeries<T, C>
319where
320 T: FormalPowerSeriesCoefficient,
321{
322 type Output = FormalPowerSeries<T, C>;
323 fn neg(self) -> Self::Output {
324 self.clone().neg()
325 }
326}
327
328impl<T, C> ShrAssign<usize> for FormalPowerSeries<T, C>
329where
330 T: FormalPowerSeriesCoefficient,
331{
332 fn shr_assign(&mut self, rhs: usize) {
333 if self.length() <= rhs {
334 *self = Self::zero();
335 } else {
336 for i in rhs..self.length() {
337 self[i - rhs] = self[i].clone();
338 }
339 self.truncate(self.length() - rhs);
340 }
341 }
342}
343impl<T, C> ShlAssign<usize> for FormalPowerSeries<T, C>
344where
345 T: FormalPowerSeriesCoefficient,
346{
347 fn shl_assign(&mut self, rhs: usize) {
348 let n = self.length();
349 self.resize(n + rhs);
350 for i in (0..n).rev() {
351 self[i + rhs] = self[i].clone();
352 }
353 for i in 0..rhs {
354 self[i] = T::zero();
355 }
356 }
357}
358
359impl<T, C> Shr<usize> for FormalPowerSeries<T, C>
360where
361 T: FormalPowerSeriesCoefficient,
362{
363 type Output = Self;
364 fn shr(mut self, rhs: usize) -> Self::Output {
365 self.shr_assign(rhs);
366 self
367 }
368}
369impl<T, C> Shl<usize> for FormalPowerSeries<T, C>
370where
371 T: FormalPowerSeriesCoefficient,
372{
373 type Output = Self;
374 fn shl(mut self, rhs: usize) -> Self::Output {
375 self.shl_assign(rhs);
376 self
377 }
378}
379impl<T, C> Shr<usize> for &FormalPowerSeries<T, C>
380where
381 T: FormalPowerSeriesCoefficient,
382{
383 type Output = FormalPowerSeries<T, C>;
384 fn shr(self, rhs: usize) -> Self::Output {
385 if self.length() <= rhs {
386 Self::Output::zero()
387 } else {
388 let mut f = Self::Output::zeros(self.length() - rhs);
389 for i in rhs..self.length() {
390 f[i - rhs] = self[i].clone();
391 }
392 f
393 }
394 }
395}
396impl<T, C> Shl<usize> for &FormalPowerSeries<T, C>
397where
398 T: FormalPowerSeriesCoefficient,
399{
400 type Output = FormalPowerSeries<T, C>;
401 fn shl(self, rhs: usize) -> Self::Output {
402 let mut f = Self::Output::zeros(self.length() + rhs);
403 for (i, x) in self.iter().cloned().enumerate().rev() {
404 f[i + rhs] = x;
405 }
406 f
407 }crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 248)
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 }
265}
266
267fn degree<T, C>(fps: &FormalPowerSeries<T, C>) -> isize {
268 fps.length() as isize - 1
269}
270
271fn add<T, C>(
272 left: FormalPowerSeries<T, C>,
273 right: FormalPowerSeries<T, C>,
274) -> FormalPowerSeries<T, C>
275where
276 T: FormalPowerSeriesCoefficient,
277{
278 (left + right).trimed()
279}
280
281fn tail<T, C>(fps: &FormalPowerSeries<T, C>, start: isize) -> FormalPowerSeries<T, C>
282where
283 T: FormalPowerSeriesCoefficient,
284{
285 let start = start.max(0) as usize;
286 if start >= fps.length() {
287 FormalPowerSeries::zero()
288 } else {
289 FormalPowerSeries::from_vec(fps.data[start..].to_vec())
290 }
291}
292
293fn coefficient<T, C>(fps: &FormalPowerSeries<T, C>, index: isize) -> T
294where
295 T: FormalPowerSeriesCoefficient,
296{
297 if index < 0 {
298 T::zero()
299 } else {
300 fps.coeff(index as usize)
301 }
302}
303
304fn brute_force<T, C>(
305 mut p: FormalPowerSeries<T, C>,
306 mut q: FormalPowerSeries<T, C>,
307 k: usize,
308) -> FpsMatrix<T, C>
309where
310 T: FormalPowerSeriesCoefficient,
311 C: NttReuse<T = Vec<T>>,
312 C::F: Clone,
313{
314 let threshold = degree(&p) - k as isize;
315 let mut matrix = FpsMatrix::identity();
316 while q.length() as isize > threshold {
317 let q_degree = q.length() - 1;
318 let mut negative_quotient = vec![T::zero(); p.length() - q.length() + 1];
319 let inverse = -T::one() / &q[q_degree];
320 for i in (0..negative_quotient.len()).rev() {
321 negative_quotient[i] = p[i + q_degree].clone() * &inverse;
322 p[i + q_degree] = T::zero();
323 for j in 0..q_degree {
324 let value = negative_quotient[i].clone() * &q[j];
325 p[i + j] += &value;
326 }
327 }
328 matrix.left_multiply_step(&negative_quotient);
329 p.truncate(q_degree);
330 p.trim_tail_zeros();
331 swap(&mut p, &mut q);
332 }
333 matrix
334}
335
336fn reduced_transform<T, C>(fps: &FormalPowerSeries<T, C>, length: usize) -> C::F
337where
338 T: FormalPowerSeriesCoefficient,
339 C: NttReuse<T = Vec<T>>,
340{
341 let mut coefficients = vec![T::zero(); length];
342 for (i, value) in fps.iter().enumerate() {
343 coefficients[i & (length - 1)] += value;
344 }
345 C::transform_ntt(coefficients, length)
346}Sourcepub fn iter_mut(&mut self) -> IterMut<'_, T> ⓘ
pub fn iter_mut(&mut self) -> IterMut<'_, T> ⓘ
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_nums.rs (line 38)
37 fn mul_assign(&mut self, rhs: T) {
38 for x in self.iter_mut() {
39 x.mul_assign(&rhs);
40 }
41 }
42}
43impl<T, C> DivAssign<T> for FormalPowerSeries<T, C>
44where
45 T: FormalPowerSeriesCoefficient,
46{
47 fn div_assign(&mut self, rhs: T) {
48 let rinv = T::one() / rhs;
49 for x in self.iter_mut() {
50 x.mul_assign(&rinv);
51 }
52 }
53}
54macro_rules! impl_fps_single_binop {
55 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
56 impl<T, C> $imp_assign<&T> for FormalPowerSeries<T, C>
57 where
58 T: FormalPowerSeriesCoefficient,
59 {
60 fn $method_assign(&mut self, rhs: &T) {
61 $imp_assign::$method_assign(self, rhs.clone());
62 }
63 }
64 impl<T, C> $imp<T> for FormalPowerSeries<T, C>
65 where
66 T: FormalPowerSeriesCoefficient,
67 {
68 type Output = Self;
69 fn $method(mut self, rhs: T) -> Self::Output {
70 $imp_assign::$method_assign(&mut self, rhs);
71 self
72 }
73 }
74 impl<T, C> $imp<&T> for FormalPowerSeries<T, C>
75 where
76 T: FormalPowerSeriesCoefficient,
77 {
78 type Output = Self;
79 fn $method(mut self, rhs: &T) -> Self::Output {
80 $imp_assign::$method_assign(&mut self, rhs);
81 self
82 }
83 }
84 impl<T, C> $imp<T> for &FormalPowerSeries<T, C>
85 where
86 T: FormalPowerSeriesCoefficient,
87 {
88 type Output = FormalPowerSeries<T, C>;
89 fn $method(self, rhs: T) -> Self::Output {
90 $imp::$method(self.clone(), rhs)
91 }
92 }
93 impl<T, C> $imp<&T> for &FormalPowerSeries<T, C>
94 where
95 T: FormalPowerSeriesCoefficient,
96 {
97 type Output = FormalPowerSeries<T, C>;
98 fn $method(self, rhs: &T) -> Self::Output {
99 $imp::$method(self.clone(), rhs)
100 }
101 }
102 };
103}
104impl_fps_single_binop!(Add, add, AddAssign, add_assign);
105impl_fps_single_binop!(Sub, sub, SubAssign, sub_assign);
106impl_fps_single_binop!(Mul, mul, MulAssign, mul_assign);
107impl_fps_single_binop!(Div, div, DivAssign, div_assign);
108
109impl<T, C> AddAssign<&Self> for FormalPowerSeries<T, C>
110where
111 T: FormalPowerSeriesCoefficient,
112{
113 fn add_assign(&mut self, rhs: &Self) {
114 if self.length() < rhs.length() {
115 self.resize(rhs.length());
116 }
117 for (x, y) in self.iter_mut().zip(rhs.iter()) {
118 x.add_assign(y);
119 }
120 }
121}
122impl<T, C> SubAssign<&Self> for FormalPowerSeries<T, C>
123where
124 T: FormalPowerSeriesCoefficient,
125{
126 fn sub_assign(&mut self, rhs: &Self) {
127 if self.length() < rhs.length() {
128 self.resize(rhs.length());
129 }
130 for (x, y) in self.iter_mut().zip(rhs.iter()) {
131 x.sub_assign(y);
132 }
133 self.trim_tail_zeros();
134 }
135}
136
137macro_rules! impl_fps_binop_addsub {
138 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
139 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
140 where
141 T: FormalPowerSeriesCoefficient,
142 {
143 fn $method_assign(&mut self, rhs: Self) {
144 $imp_assign::$method_assign(self, &rhs);
145 }
146 }
147 impl<T, C> $imp for FormalPowerSeries<T, C>
148 where
149 T: FormalPowerSeriesCoefficient,
150 {
151 type Output = Self;
152 fn $method(mut self, rhs: Self) -> Self::Output {
153 $imp_assign::$method_assign(&mut self, &rhs);
154 self
155 }
156 }
157 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
158 where
159 T: FormalPowerSeriesCoefficient,
160 {
161 type Output = Self;
162 fn $method(mut self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
163 $imp_assign::$method_assign(&mut self, rhs);
164 self
165 }
166 }
167 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
168 where
169 T: FormalPowerSeriesCoefficient,
170 {
171 type Output = FormalPowerSeries<T, C>;
172 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
173 let mut self_ = self.clone();
174 $imp_assign::$method_assign(&mut self_, &rhs);
175 self_
176 }
177 }
178 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
179 where
180 T: FormalPowerSeriesCoefficient,
181 {
182 type Output = FormalPowerSeries<T, C>;
183 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
184 let mut self_ = self.clone();
185 $imp_assign::$method_assign(&mut self_, rhs);
186 self_
187 }
188 }
189 };
190}
191impl_fps_binop_addsub!(Add, add, AddAssign, add_assign);
192impl_fps_binop_addsub!(Sub, sub, SubAssign, sub_assign);
193
194impl<T, C> Mul for FormalPowerSeries<T, C>
195where
196 C: ConvolveSteps<T = Vec<T>>,
197{
198 type Output = Self;
199 fn mul(self, rhs: Self) -> Self::Output {
200 Self::from_vec(C::convolve(self.data, rhs.data))
201 }
202}
203impl<T, C> Div for FormalPowerSeries<T, C>
204where
205 T: FormalPowerSeriesCoefficient,
206 C: ConvolveSteps<T = Vec<T>>,
207{
208 type Output = Self;
209 fn div(mut self, mut rhs: Self) -> Self::Output {
210 self.trim_tail_zeros();
211 rhs.trim_tail_zeros();
212 if self.length() < rhs.length() {
213 return Self::zero();
214 }
215 self.data.reverse();
216 rhs.data.reverse();
217 let n = self.length() - rhs.length() + 1;
218 let mut res = self * rhs.inv(n);
219 res.truncate(n);
220 res.data.reverse();
221 res
222 }
223}
224impl<T, C> Rem for FormalPowerSeries<T, C>
225where
226 T: FormalPowerSeriesCoefficient,
227 C: ConvolveSteps<T = Vec<T>>,
228{
229 type Output = Self;
230 fn rem(self, rhs: Self) -> Self::Output {
231 let mut rem = self.clone() - self / rhs.clone() * rhs;
232 rem.trim_tail_zeros();
233 rem
234 }
235}
236
237impl<T, C> FormalPowerSeries<T, C>
238where
239 T: FormalPowerSeriesCoefficient,
240 C: ConvolveSteps<T = Vec<T>>,
241{
242 pub fn div_rem(self, rhs: Self) -> (Self, Self) {
243 let div = self.clone() / rhs.clone();
244 let mut rem = self - div.clone() * rhs;
245 rem.trim_tail_zeros();
246 (div, rem)
247 }
248}
249
250macro_rules! impl_fps_binop_conv {
251 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
252 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
253 where
254 T: FormalPowerSeriesCoefficient,
255 C: ConvolveSteps<T = Vec<T>>,
256 {
257 fn $method_assign(&mut self, rhs: Self) {
258 *self = $imp::$method(Self::from_vec(take(&mut self.data)), rhs);
259 }
260 }
261 impl<T, C> $imp_assign<&Self> for FormalPowerSeries<T, C>
262 where
263 T: FormalPowerSeriesCoefficient,
264 C: ConvolveSteps<T = Vec<T>>,
265 {
266 fn $method_assign(&mut self, rhs: &Self) {
267 $imp_assign::$method_assign(self, rhs.clone());
268 }
269 }
270 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
271 where
272 T: FormalPowerSeriesCoefficient,
273 C: ConvolveSteps<T = Vec<T>>,
274 {
275 type Output = Self;
276 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
277 $imp::$method(self, rhs.clone())
278 }
279 }
280 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
281 where
282 T: FormalPowerSeriesCoefficient,
283 C: ConvolveSteps<T = Vec<T>>,
284 {
285 type Output = FormalPowerSeries<T, C>;
286 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
287 $imp::$method(self.clone(), rhs)
288 }
289 }
290 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
291 where
292 T: FormalPowerSeriesCoefficient,
293 C: ConvolveSteps<T = Vec<T>>,
294 {
295 type Output = FormalPowerSeries<T, C>;
296 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
297 $imp::$method(self.clone(), rhs.clone())
298 }
299 }
300 };
301}
302impl_fps_binop_conv!(Mul, mul, MulAssign, mul_assign);
303impl_fps_binop_conv!(Div, div, DivAssign, div_assign);
304impl_fps_binop_conv!(Rem, rem, RemAssign, rem_assign);
305
306impl<T, C> Neg for FormalPowerSeries<T, C>
307where
308 T: FormalPowerSeriesCoefficient,
309{
310 type Output = Self;
311 fn neg(mut self) -> Self::Output {
312 for x in self.iter_mut() {
313 *x = -x.clone();
314 }
315 self
316 }More examples
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 229)
228 pub fn parity_inversion(mut self) -> Self {
229 self.iter_mut()
230 .skip(1)
231 .step_by(2)
232 .for_each(|x| *x = -x.clone());
233 self
234 }
235 pub fn eval(&self, x: T) -> T {
236 self.iter()
237 .rev()
238 .fold(T::zero(), |sum, a| x.clone() * sum + a.clone())
239 }
240}
241
242impl<T, C> FormalPowerSeries<T, C>
243where
244 T: FormalPowerSeriesCoefficient,
245 C: ConvolveSteps<T = Vec<T>>,
246{
247 #[inline]
248 fn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize> {
249 let limit = deg.next_power_of_two().trailing_zeros() as usize * factor;
250 let mut count = 0;
251 let mut step = 0;
252 for (i, value) in self.iter().take(deg).enumerate() {
253 if value.is_zero() {
254 continue;
255 }
256 count += 1;
257 if step != 1 {
258 step = gcd(step, i as u64);
259 }
260 if count > limit {
261 return None;
262 }
263 }
264 Some(step.max(1) as usize)
265 }
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 256)
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 }Source§impl<T, C> FormalPowerSeries<T, C>where
T: Zero,
impl<T, C> FormalPowerSeries<T, C>where
T: Zero,
Sourcepub fn zeros(deg: usize) -> Self
pub fn zeros(deg: usize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_nums.rs (line 388)
384 fn shr(self, rhs: usize) -> Self::Output {
385 if self.length() <= rhs {
386 Self::Output::zero()
387 } else {
388 let mut f = Self::Output::zeros(self.length() - rhs);
389 for i in rhs..self.length() {
390 f[i - rhs] = self[i].clone();
391 }
392 f
393 }
394 }
395}
396impl<T, C> Shl<usize> for &FormalPowerSeries<T, C>
397where
398 T: FormalPowerSeriesCoefficient,
399{
400 type Output = FormalPowerSeries<T, C>;
401 fn shl(self, rhs: usize) -> Self::Output {
402 let mut f = Self::Output::zeros(self.length() + rhs);
403 for (i, x) in self.iter().cloned().enumerate().rev() {
404 f[i + rhs] = x;
405 }
406 f
407 }More examples
crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 238)
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 }
265}
266
267fn degree<T, C>(fps: &FormalPowerSeries<T, C>) -> isize {
268 fps.length() as isize - 1
269}
270
271fn add<T, C>(
272 left: FormalPowerSeries<T, C>,
273 right: FormalPowerSeries<T, C>,
274) -> FormalPowerSeries<T, C>
275where
276 T: FormalPowerSeriesCoefficient,
277{
278 (left + right).trimed()
279}
280
281fn tail<T, C>(fps: &FormalPowerSeries<T, C>, start: isize) -> FormalPowerSeries<T, C>
282where
283 T: FormalPowerSeriesCoefficient,
284{
285 let start = start.max(0) as usize;
286 if start >= fps.length() {
287 FormalPowerSeries::zero()
288 } else {
289 FormalPowerSeries::from_vec(fps.data[start..].to_vec())
290 }
291}
292
293fn coefficient<T, C>(fps: &FormalPowerSeries<T, C>, index: isize) -> T
294where
295 T: FormalPowerSeriesCoefficient,
296{
297 if index < 0 {
298 T::zero()
299 } else {
300 fps.coeff(index as usize)
301 }
302}
303
304fn brute_force<T, C>(
305 mut p: FormalPowerSeries<T, C>,
306 mut q: FormalPowerSeries<T, C>,
307 k: usize,
308) -> FpsMatrix<T, C>
309where
310 T: FormalPowerSeriesCoefficient,
311 C: NttReuse<T = Vec<T>>,
312 C::F: Clone,
313{
314 let threshold = degree(&p) - k as isize;
315 let mut matrix = FpsMatrix::identity();
316 while q.length() as isize > threshold {
317 let q_degree = q.length() - 1;
318 let mut negative_quotient = vec![T::zero(); p.length() - q.length() + 1];
319 let inverse = -T::one() / &q[q_degree];
320 for i in (0..negative_quotient.len()).rev() {
321 negative_quotient[i] = p[i + q_degree].clone() * &inverse;
322 p[i + q_degree] = T::zero();
323 for j in 0..q_degree {
324 let value = negative_quotient[i].clone() * &q[j];
325 p[i + j] += &value;
326 }
327 }
328 matrix.left_multiply_step(&negative_quotient);
329 p.truncate(q_degree);
330 p.trim_tail_zeros();
331 swap(&mut p, &mut q);
332 }
333 matrix
334}
335
336fn reduced_transform<T, C>(fps: &FormalPowerSeries<T, C>, length: usize) -> C::F
337where
338 T: FormalPowerSeriesCoefficient,
339 C: NttReuse<T = Vec<T>>,
340{
341 let mut coefficients = vec![T::zero(); length];
342 for (i, value) in fps.iter().enumerate() {
343 coefficients[i & (length - 1)] += value;
344 }
345 C::transform_ntt(coefficients, length)
346}
347
348fn transform_window<T, C>(fps: &FormalPowerSeries<T, C>, end: isize, length: usize) -> C::F
349where
350 T: FormalPowerSeriesCoefficient,
351 C: NttReuse<T = Vec<T>>,
352{
353 let start = end - length as isize;
354 let coefficients = (start..end).map(|index| coefficient(fps, index)).collect();
355 C::transform_ntt(coefficients, length)
356}
357
358fn half_gcd<T, C>(
359 p: &FormalPowerSeries<T, C>,
360 q: &FormalPowerSeries<T, C>,
361 k: usize,
362 length: usize,
363) -> (FpsMatrix<T, C>, FrequencyMatrix<C>)
364where
365 T: FormalPowerSeriesCoefficient,
366 C: NttReuse<T = Vec<T>>,
367 C::F: Clone,
368{
369 let d = degree(p);
370 if degree(q) < d - k as isize {
371 let matrix = FpsMatrix::identity();
372 let frequency = matrix.transform(length);
373 return (matrix, frequency);
374 }
375 if k == 1 {
376 let matrix = FpsMatrix {
377 a00: FormalPowerSeries::zero(),
378 a01: FormalPowerSeries::one(),
379 a10: FormalPowerSeries::one(),
380 a11: -(tail(p, d - 2) / tail(q, d - 2)),
381 };
382 let frequency = matrix.transform(length);
383 return (matrix, frequency);
384 }
385 if p.length().min(q.length()) <= 32 {
386 let matrix = brute_force(p.clone(), q.clone(), k);
387 let frequency = matrix.transform(length);
388 return (matrix, frequency);
389 }
390
391 let half = length / 2;
392 if k <= half {
393 let (matrix, frequency) = half_gcd(p, q, k, half);
394 let frequency = matrix.extend_transform(frequency, length);
395 return (matrix, frequency);
396 }
397
398 let (matrix, mut matrix_frequency) = half_gcd(
399 &tail(p, d - 2 * half as isize),
400 &tail(q, d - 2 * half as isize),
401 half,
402 length,
403 );
404 let degeneracy = half as isize - degree(&matrix.a11);
405
406 let (p0, q0) = matrix_frequency.apply(
407 &transform_window(p, d - half as isize + degeneracy, length),
408 &transform_window(q, d - half as isize + degeneracy, length),
409 length,
410 );
411 let (p1, q1) = matrix_frequency.apply(
412 &transform_window(p, d - 2 * half as isize, length),
413 &transform_window(q, d - 2 * half as isize, length),
414 length,
415 );
416 let part_length = (half as isize + degeneracy) as usize;
417 let mut p_reduced = p1[length - part_length..].to_vec();
418 p_reduced.extend_from_slice(&p0[length - part_length..]);
419 let mut q_reduced = q1[length - part_length..].to_vec();
420 q_reduced.extend_from_slice(&q0[length - part_length..]);
421 let mut q_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(q_reduced).trimed();
422
423 let position = d - half as isize + degeneracy;
424 let mut leading = T::zero();
425 for i in 0..=position {
426 leading += coefficient(p, i) * coefficient(&matrix.a00, position - i)
427 + coefficient(q, i) * coefficient(&matrix.a01, position - i);
428 }
429 p_reduced.push(leading);
430 let mut p_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(p_reduced);
431 if degree(&q_reduced) < 3 * half as isize + degeneracy - k as isize {
432 return (matrix, matrix_frequency);
433 }
434
435 let mut remaining = k as isize - degree(&matrix.a11);
436 let mut top_product = matrix.a11.data.last().unwrap().clone();
437 let mut product_degree = degree(&matrix.a11);
438 if degeneracy > 0 {
439 let skip = (2 * half as isize + 2 * degeneracy - (d - half as isize + degeneracy)).max(0);
440 let (division, remainder) = tail(&p_reduced, skip).div_rem(tail(&q_reduced, skip));
441 remaining -= degree(&division);
442 top_product *= -division.data.last().unwrap().clone();
443 product_degree += degree(&division);
444 matrix_frequency = matrix_frequency.left_multiply_step(&division, length);
445 swap(&mut p_reduced, &mut q_reduced);
446 q_reduced = FormalPowerSeries::zeros(skip as usize);
447 q_reduced.data.extend(remainder.data);
448 }
449
450 let start = 3 * half as isize + degeneracy - k as isize - remaining;
451 let (right_matrix, right_frequency) = half_gcd(
452 &tail(&p_reduced, start),
453 &tail(&q_reduced, start),
454 remaining as usize,
455 length,
456 );
457 let product_frequency = right_frequency.multiply(&matrix_frequency);
458 let mut product = product_frequency.clone().inverse_transform(length);
459 product.a00.truncate(k);
460 product.a00.trim_tail_zeros();
461 product.a01.truncate(k);
462 product.a01.trim_tail_zeros();
463 product.a10.truncate(k);
464 product.a10.trim_tail_zeros();
465 product_degree += degree(&right_matrix.a11);
466 if product_degree == length as isize {
467 product.a11.resize(k + 1);
468 let highest = top_product * right_matrix.a11.data.last().unwrap();
469 product.a11[k] = highest.clone();
470 product.a11[0] -= highest;
471 }
472 product.a11.trim_tail_zeros();
473 let product_frequency = if C::MULTIPLE {
474 product.transform(length)
475 } else {
476 product_frequency
477 };
478 (product, product_frequency)
479}crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 282)
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }Sourcepub fn resize(&mut self, deg: usize)
pub fn resize(&mut self, deg: usize)
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 62)
61 pub fn resized(mut self, deg: usize) -> Self {
62 self.resize(deg);
63 self
64 }
65 pub fn reversed(mut self) -> Self {
66 self.data.reverse();
67 self
68 }
69}
70
71impl<T, C> FormalPowerSeries<T, C>
72where
73 T: Zero + Clone,
74{
75 pub fn coeff(&self, deg: usize) -> T {
76 self.data.get(deg).cloned().unwrap_or_else(T::zero)
77 }
78}
79
80impl<T, C> FormalPowerSeries<T, C>
81where
82 T: Zero + PartialEq,
83{
84 pub fn trim_tail_zeros(&mut self) {
85 let mut len = self.length();
86 while len > 0 {
87 if self.data[len - 1].is_zero() {
88 len -= 1;
89 } else {
90 break;
91 }
92 }
93 self.truncate(len);
94 }
95 pub fn trimed(mut self) -> Self {
96 self.trim_tail_zeros();
97 self
98 }
99}
100
101impl<T, C> Zero for FormalPowerSeries<T, C>
102where
103 T: PartialEq,
104{
105 fn zero() -> Self {
106 Self::from_vec(Vec::new())
107 }
108}
109impl<T, C> One for FormalPowerSeries<T, C>
110where
111 T: PartialEq + One,
112{
113 fn one() -> Self {
114 Self::from(T::one())
115 }
116}
117
118impl<T, C> IntoIterator for FormalPowerSeries<T, C> {
119 type Item = T;
120 type IntoIter = std::vec::IntoIter<T>;
121 fn into_iter(self) -> Self::IntoIter {
122 self.data.into_iter()
123 }
124}
125impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C> {
126 type Item = &'a T;
127 type IntoIter = Iter<'a, T>;
128 fn into_iter(self) -> Self::IntoIter {
129 self.data.iter()
130 }
131}
132impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C> {
133 type Item = &'a mut T;
134 type IntoIter = IterMut<'a, T>;
135 fn into_iter(self) -> Self::IntoIter {
136 self.data.iter_mut()
137 }
138}
139
140impl<T, C> FromIterator<T> for FormalPowerSeries<T, C> {
141 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
142 Self::from_vec(iter.into_iter().collect())
143 }
144}
145
146impl<T, C> Index<usize> for FormalPowerSeries<T, C> {
147 type Output = T;
148 fn index(&self, index: usize) -> &Self::Output {
149 &self.data[index]
150 }
151}
152impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C> {
153 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
154 &mut self.data[index]
155 }
156}
157
158impl<T, C> From<T> for FormalPowerSeries<T, C> {
159 fn from(x: T) -> Self {
160 once(x).collect()
161 }
162}
163impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C> {
164 fn from(data: Vec<T>) -> Self {
165 Self::from_vec(data)
166 }
167}
168
169impl<T, C> FormalPowerSeries<T, C>
170where
171 T: FormalPowerSeriesCoefficient,
172{
173 pub fn prefix_ref(&self, deg: usize) -> Self {
174 if deg < self.length() {
175 Self::from_vec(self.data[..deg].to_vec())
176 } else {
177 self.clone()
178 }
179 }
180 pub fn prefix(mut self, deg: usize) -> Self {
181 self.data.truncate(deg);
182 self
183 }
184 pub fn even(mut self) -> Self {
185 let mut keep = false;
186 self.data.retain(|_| {
187 keep = !keep;
188 keep
189 });
190 self
191 }
192 pub fn odd(mut self) -> Self {
193 let mut keep = true;
194 self.data.retain(|_| {
195 keep = !keep;
196 keep
197 });
198 self
199 }
200 pub fn diff(mut self) -> Self {
201 let mut c = T::one();
202 for i in 1..self.length() {
203 self.data[i - 1] = self.data[i].clone() * &c;
204 c += T::one();
205 }
206 self.data.pop();
207 self
208 }
209 pub fn integral(mut self) -> Self {
210 let n = self.length();
211 let mut fact = Vec::with_capacity(n + 1);
212 let mut c = T::one();
213 fact.push(c.clone());
214 for _ in 1..n {
215 fact.push(fact.last().cloned().unwrap() * c.clone());
216 c += T::one();
217 }
218 let mut invf = T::one() / (fact.last().cloned().unwrap() * c.clone());
219 self.data.push(T::zero());
220 for i in (1..=n).rev() {
221 self.data[i] = self.data[i - 1].clone() * (invf.clone() * fact.pop().unwrap());
222 invf *= c.clone();
223 c -= T::one();
224 }
225 self.data[0] = T::zero();
226 self
227 }
228 pub fn parity_inversion(mut self) -> Self {
229 self.iter_mut()
230 .skip(1)
231 .step_by(2)
232 .for_each(|x| *x = -x.clone());
233 self
234 }
235 pub fn eval(&self, x: T) -> T {
236 self.iter()
237 .rev()
238 .fold(T::zero(), |sum, a| x.clone() * sum + a.clone())
239 }
240}
241
242impl<T, C> FormalPowerSeries<T, C>
243where
244 T: FormalPowerSeriesCoefficient,
245 C: ConvolveSteps<T = Vec<T>>,
246{
247 #[inline]
248 fn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize> {
249 let limit = deg.next_power_of_two().trailing_zeros() as usize * factor;
250 let mut count = 0;
251 let mut step = 0;
252 for (i, value) in self.iter().take(deg).enumerate() {
253 if value.is_zero() {
254 continue;
255 }
256 count += 1;
257 if step != 1 {
258 step = gcd(step, i as u64);
259 }
260 if count > limit {
261 return None;
262 }
263 }
264 Some(step.max(1) as usize)
265 }
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }More examples
crates/competitive/src/math/formal_power_series/formal_power_series_nums.rs (line 115)
113 fn add_assign(&mut self, rhs: &Self) {
114 if self.length() < rhs.length() {
115 self.resize(rhs.length());
116 }
117 for (x, y) in self.iter_mut().zip(rhs.iter()) {
118 x.add_assign(y);
119 }
120 }
121}
122impl<T, C> SubAssign<&Self> for FormalPowerSeries<T, C>
123where
124 T: FormalPowerSeriesCoefficient,
125{
126 fn sub_assign(&mut self, rhs: &Self) {
127 if self.length() < rhs.length() {
128 self.resize(rhs.length());
129 }
130 for (x, y) in self.iter_mut().zip(rhs.iter()) {
131 x.sub_assign(y);
132 }
133 self.trim_tail_zeros();
134 }
135}
136
137macro_rules! impl_fps_binop_addsub {
138 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
139 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
140 where
141 T: FormalPowerSeriesCoefficient,
142 {
143 fn $method_assign(&mut self, rhs: Self) {
144 $imp_assign::$method_assign(self, &rhs);
145 }
146 }
147 impl<T, C> $imp for FormalPowerSeries<T, C>
148 where
149 T: FormalPowerSeriesCoefficient,
150 {
151 type Output = Self;
152 fn $method(mut self, rhs: Self) -> Self::Output {
153 $imp_assign::$method_assign(&mut self, &rhs);
154 self
155 }
156 }
157 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
158 where
159 T: FormalPowerSeriesCoefficient,
160 {
161 type Output = Self;
162 fn $method(mut self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
163 $imp_assign::$method_assign(&mut self, rhs);
164 self
165 }
166 }
167 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
168 where
169 T: FormalPowerSeriesCoefficient,
170 {
171 type Output = FormalPowerSeries<T, C>;
172 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
173 let mut self_ = self.clone();
174 $imp_assign::$method_assign(&mut self_, &rhs);
175 self_
176 }
177 }
178 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
179 where
180 T: FormalPowerSeriesCoefficient,
181 {
182 type Output = FormalPowerSeries<T, C>;
183 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
184 let mut self_ = self.clone();
185 $imp_assign::$method_assign(&mut self_, rhs);
186 self_
187 }
188 }
189 };
190}
191impl_fps_binop_addsub!(Add, add, AddAssign, add_assign);
192impl_fps_binop_addsub!(Sub, sub, SubAssign, sub_assign);
193
194impl<T, C> Mul for FormalPowerSeries<T, C>
195where
196 C: ConvolveSteps<T = Vec<T>>,
197{
198 type Output = Self;
199 fn mul(self, rhs: Self) -> Self::Output {
200 Self::from_vec(C::convolve(self.data, rhs.data))
201 }
202}
203impl<T, C> Div for FormalPowerSeries<T, C>
204where
205 T: FormalPowerSeriesCoefficient,
206 C: ConvolveSteps<T = Vec<T>>,
207{
208 type Output = Self;
209 fn div(mut self, mut rhs: Self) -> Self::Output {
210 self.trim_tail_zeros();
211 rhs.trim_tail_zeros();
212 if self.length() < rhs.length() {
213 return Self::zero();
214 }
215 self.data.reverse();
216 rhs.data.reverse();
217 let n = self.length() - rhs.length() + 1;
218 let mut res = self * rhs.inv(n);
219 res.truncate(n);
220 res.data.reverse();
221 res
222 }
223}
224impl<T, C> Rem for FormalPowerSeries<T, C>
225where
226 T: FormalPowerSeriesCoefficient,
227 C: ConvolveSteps<T = Vec<T>>,
228{
229 type Output = Self;
230 fn rem(self, rhs: Self) -> Self::Output {
231 let mut rem = self.clone() - self / rhs.clone() * rhs;
232 rem.trim_tail_zeros();
233 rem
234 }
235}
236
237impl<T, C> FormalPowerSeries<T, C>
238where
239 T: FormalPowerSeriesCoefficient,
240 C: ConvolveSteps<T = Vec<T>>,
241{
242 pub fn div_rem(self, rhs: Self) -> (Self, Self) {
243 let div = self.clone() / rhs.clone();
244 let mut rem = self - div.clone() * rhs;
245 rem.trim_tail_zeros();
246 (div, rem)
247 }
248}
249
250macro_rules! impl_fps_binop_conv {
251 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
252 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
253 where
254 T: FormalPowerSeriesCoefficient,
255 C: ConvolveSteps<T = Vec<T>>,
256 {
257 fn $method_assign(&mut self, rhs: Self) {
258 *self = $imp::$method(Self::from_vec(take(&mut self.data)), rhs);
259 }
260 }
261 impl<T, C> $imp_assign<&Self> for FormalPowerSeries<T, C>
262 where
263 T: FormalPowerSeriesCoefficient,
264 C: ConvolveSteps<T = Vec<T>>,
265 {
266 fn $method_assign(&mut self, rhs: &Self) {
267 $imp_assign::$method_assign(self, rhs.clone());
268 }
269 }
270 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
271 where
272 T: FormalPowerSeriesCoefficient,
273 C: ConvolveSteps<T = Vec<T>>,
274 {
275 type Output = Self;
276 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
277 $imp::$method(self, rhs.clone())
278 }
279 }
280 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
281 where
282 T: FormalPowerSeriesCoefficient,
283 C: ConvolveSteps<T = Vec<T>>,
284 {
285 type Output = FormalPowerSeries<T, C>;
286 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
287 $imp::$method(self.clone(), rhs)
288 }
289 }
290 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
291 where
292 T: FormalPowerSeriesCoefficient,
293 C: ConvolveSteps<T = Vec<T>>,
294 {
295 type Output = FormalPowerSeries<T, C>;
296 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
297 $imp::$method(self.clone(), rhs.clone())
298 }
299 }
300 };
301}
302impl_fps_binop_conv!(Mul, mul, MulAssign, mul_assign);
303impl_fps_binop_conv!(Div, div, DivAssign, div_assign);
304impl_fps_binop_conv!(Rem, rem, RemAssign, rem_assign);
305
306impl<T, C> Neg for FormalPowerSeries<T, C>
307where
308 T: FormalPowerSeriesCoefficient,
309{
310 type Output = Self;
311 fn neg(mut self) -> Self::Output {
312 for x in self.iter_mut() {
313 *x = -x.clone();
314 }
315 self
316 }
317}
318impl<T, C> Neg for &FormalPowerSeries<T, C>
319where
320 T: FormalPowerSeriesCoefficient,
321{
322 type Output = FormalPowerSeries<T, C>;
323 fn neg(self) -> Self::Output {
324 self.clone().neg()
325 }
326}
327
328impl<T, C> ShrAssign<usize> for FormalPowerSeries<T, C>
329where
330 T: FormalPowerSeriesCoefficient,
331{
332 fn shr_assign(&mut self, rhs: usize) {
333 if self.length() <= rhs {
334 *self = Self::zero();
335 } else {
336 for i in rhs..self.length() {
337 self[i - rhs] = self[i].clone();
338 }
339 self.truncate(self.length() - rhs);
340 }
341 }
342}
343impl<T, C> ShlAssign<usize> for FormalPowerSeries<T, C>
344where
345 T: FormalPowerSeriesCoefficient,
346{
347 fn shl_assign(&mut self, rhs: usize) {
348 let n = self.length();
349 self.resize(n + rhs);
350 for i in (0..n).rev() {
351 self[i + rhs] = self[i].clone();
352 }
353 for i in 0..rhs {
354 self[i] = T::zero();
355 }
356 }crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 261)
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 }
265}
266
267fn degree<T, C>(fps: &FormalPowerSeries<T, C>) -> isize {
268 fps.length() as isize - 1
269}
270
271fn add<T, C>(
272 left: FormalPowerSeries<T, C>,
273 right: FormalPowerSeries<T, C>,
274) -> FormalPowerSeries<T, C>
275where
276 T: FormalPowerSeriesCoefficient,
277{
278 (left + right).trimed()
279}
280
281fn tail<T, C>(fps: &FormalPowerSeries<T, C>, start: isize) -> FormalPowerSeries<T, C>
282where
283 T: FormalPowerSeriesCoefficient,
284{
285 let start = start.max(0) as usize;
286 if start >= fps.length() {
287 FormalPowerSeries::zero()
288 } else {
289 FormalPowerSeries::from_vec(fps.data[start..].to_vec())
290 }
291}
292
293fn coefficient<T, C>(fps: &FormalPowerSeries<T, C>, index: isize) -> T
294where
295 T: FormalPowerSeriesCoefficient,
296{
297 if index < 0 {
298 T::zero()
299 } else {
300 fps.coeff(index as usize)
301 }
302}
303
304fn brute_force<T, C>(
305 mut p: FormalPowerSeries<T, C>,
306 mut q: FormalPowerSeries<T, C>,
307 k: usize,
308) -> FpsMatrix<T, C>
309where
310 T: FormalPowerSeriesCoefficient,
311 C: NttReuse<T = Vec<T>>,
312 C::F: Clone,
313{
314 let threshold = degree(&p) - k as isize;
315 let mut matrix = FpsMatrix::identity();
316 while q.length() as isize > threshold {
317 let q_degree = q.length() - 1;
318 let mut negative_quotient = vec![T::zero(); p.length() - q.length() + 1];
319 let inverse = -T::one() / &q[q_degree];
320 for i in (0..negative_quotient.len()).rev() {
321 negative_quotient[i] = p[i + q_degree].clone() * &inverse;
322 p[i + q_degree] = T::zero();
323 for j in 0..q_degree {
324 let value = negative_quotient[i].clone() * &q[j];
325 p[i + j] += &value;
326 }
327 }
328 matrix.left_multiply_step(&negative_quotient);
329 p.truncate(q_degree);
330 p.trim_tail_zeros();
331 swap(&mut p, &mut q);
332 }
333 matrix
334}
335
336fn reduced_transform<T, C>(fps: &FormalPowerSeries<T, C>, length: usize) -> C::F
337where
338 T: FormalPowerSeriesCoefficient,
339 C: NttReuse<T = Vec<T>>,
340{
341 let mut coefficients = vec![T::zero(); length];
342 for (i, value) in fps.iter().enumerate() {
343 coefficients[i & (length - 1)] += value;
344 }
345 C::transform_ntt(coefficients, length)
346}
347
348fn transform_window<T, C>(fps: &FormalPowerSeries<T, C>, end: isize, length: usize) -> C::F
349where
350 T: FormalPowerSeriesCoefficient,
351 C: NttReuse<T = Vec<T>>,
352{
353 let start = end - length as isize;
354 let coefficients = (start..end).map(|index| coefficient(fps, index)).collect();
355 C::transform_ntt(coefficients, length)
356}
357
358fn half_gcd<T, C>(
359 p: &FormalPowerSeries<T, C>,
360 q: &FormalPowerSeries<T, C>,
361 k: usize,
362 length: usize,
363) -> (FpsMatrix<T, C>, FrequencyMatrix<C>)
364where
365 T: FormalPowerSeriesCoefficient,
366 C: NttReuse<T = Vec<T>>,
367 C::F: Clone,
368{
369 let d = degree(p);
370 if degree(q) < d - k as isize {
371 let matrix = FpsMatrix::identity();
372 let frequency = matrix.transform(length);
373 return (matrix, frequency);
374 }
375 if k == 1 {
376 let matrix = FpsMatrix {
377 a00: FormalPowerSeries::zero(),
378 a01: FormalPowerSeries::one(),
379 a10: FormalPowerSeries::one(),
380 a11: -(tail(p, d - 2) / tail(q, d - 2)),
381 };
382 let frequency = matrix.transform(length);
383 return (matrix, frequency);
384 }
385 if p.length().min(q.length()) <= 32 {
386 let matrix = brute_force(p.clone(), q.clone(), k);
387 let frequency = matrix.transform(length);
388 return (matrix, frequency);
389 }
390
391 let half = length / 2;
392 if k <= half {
393 let (matrix, frequency) = half_gcd(p, q, k, half);
394 let frequency = matrix.extend_transform(frequency, length);
395 return (matrix, frequency);
396 }
397
398 let (matrix, mut matrix_frequency) = half_gcd(
399 &tail(p, d - 2 * half as isize),
400 &tail(q, d - 2 * half as isize),
401 half,
402 length,
403 );
404 let degeneracy = half as isize - degree(&matrix.a11);
405
406 let (p0, q0) = matrix_frequency.apply(
407 &transform_window(p, d - half as isize + degeneracy, length),
408 &transform_window(q, d - half as isize + degeneracy, length),
409 length,
410 );
411 let (p1, q1) = matrix_frequency.apply(
412 &transform_window(p, d - 2 * half as isize, length),
413 &transform_window(q, d - 2 * half as isize, length),
414 length,
415 );
416 let part_length = (half as isize + degeneracy) as usize;
417 let mut p_reduced = p1[length - part_length..].to_vec();
418 p_reduced.extend_from_slice(&p0[length - part_length..]);
419 let mut q_reduced = q1[length - part_length..].to_vec();
420 q_reduced.extend_from_slice(&q0[length - part_length..]);
421 let mut q_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(q_reduced).trimed();
422
423 let position = d - half as isize + degeneracy;
424 let mut leading = T::zero();
425 for i in 0..=position {
426 leading += coefficient(p, i) * coefficient(&matrix.a00, position - i)
427 + coefficient(q, i) * coefficient(&matrix.a01, position - i);
428 }
429 p_reduced.push(leading);
430 let mut p_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(p_reduced);
431 if degree(&q_reduced) < 3 * half as isize + degeneracy - k as isize {
432 return (matrix, matrix_frequency);
433 }
434
435 let mut remaining = k as isize - degree(&matrix.a11);
436 let mut top_product = matrix.a11.data.last().unwrap().clone();
437 let mut product_degree = degree(&matrix.a11);
438 if degeneracy > 0 {
439 let skip = (2 * half as isize + 2 * degeneracy - (d - half as isize + degeneracy)).max(0);
440 let (division, remainder) = tail(&p_reduced, skip).div_rem(tail(&q_reduced, skip));
441 remaining -= degree(&division);
442 top_product *= -division.data.last().unwrap().clone();
443 product_degree += degree(&division);
444 matrix_frequency = matrix_frequency.left_multiply_step(&division, length);
445 swap(&mut p_reduced, &mut q_reduced);
446 q_reduced = FormalPowerSeries::zeros(skip as usize);
447 q_reduced.data.extend(remainder.data);
448 }
449
450 let start = 3 * half as isize + degeneracy - k as isize - remaining;
451 let (right_matrix, right_frequency) = half_gcd(
452 &tail(&p_reduced, start),
453 &tail(&q_reduced, start),
454 remaining as usize,
455 length,
456 );
457 let product_frequency = right_frequency.multiply(&matrix_frequency);
458 let mut product = product_frequency.clone().inverse_transform(length);
459 product.a00.truncate(k);
460 product.a00.trim_tail_zeros();
461 product.a01.truncate(k);
462 product.a01.trim_tail_zeros();
463 product.a10.truncate(k);
464 product.a10.trim_tail_zeros();
465 product_degree += degree(&right_matrix.a11);
466 if product_degree == length as isize {
467 product.a11.resize(k + 1);
468 let highest = top_product * right_matrix.a11.data.last().unwrap();
469 product.a11[k] = highest.clone();
470 product.a11[0] -= highest;
471 }
472 product.a11.trim_tail_zeros();
473 let product_frequency = if C::MULTIPLE {
474 product.transform(length)
475 } else {
476 product_frequency
477 };
478 (product, product_frequency)
479}Sourcepub fn resized(self, deg: usize) -> Self
pub fn resized(self, deg: usize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 659)
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }Sourcepub fn reversed(self) -> Self
pub fn reversed(self) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 1071)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }Source§impl<T, C> FormalPowerSeries<T, C>
impl<T, C> FormalPowerSeries<T, C>
Sourcepub fn coeff(&self, deg: usize) -> T
pub fn coeff(&self, deg: usize) -> T
Examples found in repository?
More examples
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 508)
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }Source§impl<T, C> FormalPowerSeries<T, C>
impl<T, C> FormalPowerSeries<T, C>
Sourcepub fn trim_tail_zeros(&mut self)
pub fn trim_tail_zeros(&mut self)
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 96)
95 pub fn trimed(mut self) -> Self {
96 self.trim_tail_zeros();
97 self
98 }
99}
100
101impl<T, C> Zero for FormalPowerSeries<T, C>
102where
103 T: PartialEq,
104{
105 fn zero() -> Self {
106 Self::from_vec(Vec::new())
107 }
108}
109impl<T, C> One for FormalPowerSeries<T, C>
110where
111 T: PartialEq + One,
112{
113 fn one() -> Self {
114 Self::from(T::one())
115 }
116}
117
118impl<T, C> IntoIterator for FormalPowerSeries<T, C> {
119 type Item = T;
120 type IntoIter = std::vec::IntoIter<T>;
121 fn into_iter(self) -> Self::IntoIter {
122 self.data.into_iter()
123 }
124}
125impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C> {
126 type Item = &'a T;
127 type IntoIter = Iter<'a, T>;
128 fn into_iter(self) -> Self::IntoIter {
129 self.data.iter()
130 }
131}
132impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C> {
133 type Item = &'a mut T;
134 type IntoIter = IterMut<'a, T>;
135 fn into_iter(self) -> Self::IntoIter {
136 self.data.iter_mut()
137 }
138}
139
140impl<T, C> FromIterator<T> for FormalPowerSeries<T, C> {
141 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
142 Self::from_vec(iter.into_iter().collect())
143 }
144}
145
146impl<T, C> Index<usize> for FormalPowerSeries<T, C> {
147 type Output = T;
148 fn index(&self, index: usize) -> &Self::Output {
149 &self.data[index]
150 }
151}
152impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C> {
153 fn index_mut(&mut self, index: usize) -> &mut Self::Output {
154 &mut self.data[index]
155 }
156}
157
158impl<T, C> From<T> for FormalPowerSeries<T, C> {
159 fn from(x: T) -> Self {
160 once(x).collect()
161 }
162}
163impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C> {
164 fn from(data: Vec<T>) -> Self {
165 Self::from_vec(data)
166 }
167}
168
169impl<T, C> FormalPowerSeries<T, C>
170where
171 T: FormalPowerSeriesCoefficient,
172{
173 pub fn prefix_ref(&self, deg: usize) -> Self {
174 if deg < self.length() {
175 Self::from_vec(self.data[..deg].to_vec())
176 } else {
177 self.clone()
178 }
179 }
180 pub fn prefix(mut self, deg: usize) -> Self {
181 self.data.truncate(deg);
182 self
183 }
184 pub fn even(mut self) -> Self {
185 let mut keep = false;
186 self.data.retain(|_| {
187 keep = !keep;
188 keep
189 });
190 self
191 }
192 pub fn odd(mut self) -> Self {
193 let mut keep = true;
194 self.data.retain(|_| {
195 keep = !keep;
196 keep
197 });
198 self
199 }
200 pub fn diff(mut self) -> Self {
201 let mut c = T::one();
202 for i in 1..self.length() {
203 self.data[i - 1] = self.data[i].clone() * &c;
204 c += T::one();
205 }
206 self.data.pop();
207 self
208 }
209 pub fn integral(mut self) -> Self {
210 let n = self.length();
211 let mut fact = Vec::with_capacity(n + 1);
212 let mut c = T::one();
213 fact.push(c.clone());
214 for _ in 1..n {
215 fact.push(fact.last().cloned().unwrap() * c.clone());
216 c += T::one();
217 }
218 let mut invf = T::one() / (fact.last().cloned().unwrap() * c.clone());
219 self.data.push(T::zero());
220 for i in (1..=n).rev() {
221 self.data[i] = self.data[i - 1].clone() * (invf.clone() * fact.pop().unwrap());
222 invf *= c.clone();
223 c -= T::one();
224 }
225 self.data[0] = T::zero();
226 self
227 }
228 pub fn parity_inversion(mut self) -> Self {
229 self.iter_mut()
230 .skip(1)
231 .step_by(2)
232 .for_each(|x| *x = -x.clone());
233 self
234 }
235 pub fn eval(&self, x: T) -> T {
236 self.iter()
237 .rev()
238 .fold(T::zero(), |sum, a| x.clone() * sum + a.clone())
239 }
240}
241
242impl<T, C> FormalPowerSeries<T, C>
243where
244 T: FormalPowerSeriesCoefficient,
245 C: ConvolveSteps<T = Vec<T>>,
246{
247 #[inline]
248 fn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize> {
249 let limit = deg.next_power_of_two().trailing_zeros() as usize * factor;
250 let mut count = 0;
251 let mut step = 0;
252 for (i, value) in self.iter().take(deg).enumerate() {
253 if value.is_zero() {
254 continue;
255 }
256 count += 1;
257 if step != 1 {
258 step = gcd(step, i as u64);
259 }
260 if count > limit {
261 return None;
262 }
263 }
264 Some(step.max(1) as usize)
265 }
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }More examples
crates/competitive/src/math/formal_power_series/formal_power_series_nums.rs (line 30)
25 fn sub_assign(&mut self, rhs: T) {
26 if self.length() == 0 {
27 self.data.push(T::zero());
28 }
29 self.data[0].sub_assign(rhs);
30 self.trim_tail_zeros();
31 }
32}
33impl<T, C> MulAssign<T> for FormalPowerSeries<T, C>
34where
35 T: FormalPowerSeriesCoefficient,
36{
37 fn mul_assign(&mut self, rhs: T) {
38 for x in self.iter_mut() {
39 x.mul_assign(&rhs);
40 }
41 }
42}
43impl<T, C> DivAssign<T> for FormalPowerSeries<T, C>
44where
45 T: FormalPowerSeriesCoefficient,
46{
47 fn div_assign(&mut self, rhs: T) {
48 let rinv = T::one() / rhs;
49 for x in self.iter_mut() {
50 x.mul_assign(&rinv);
51 }
52 }
53}
54macro_rules! impl_fps_single_binop {
55 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
56 impl<T, C> $imp_assign<&T> for FormalPowerSeries<T, C>
57 where
58 T: FormalPowerSeriesCoefficient,
59 {
60 fn $method_assign(&mut self, rhs: &T) {
61 $imp_assign::$method_assign(self, rhs.clone());
62 }
63 }
64 impl<T, C> $imp<T> for FormalPowerSeries<T, C>
65 where
66 T: FormalPowerSeriesCoefficient,
67 {
68 type Output = Self;
69 fn $method(mut self, rhs: T) -> Self::Output {
70 $imp_assign::$method_assign(&mut self, rhs);
71 self
72 }
73 }
74 impl<T, C> $imp<&T> for FormalPowerSeries<T, C>
75 where
76 T: FormalPowerSeriesCoefficient,
77 {
78 type Output = Self;
79 fn $method(mut self, rhs: &T) -> Self::Output {
80 $imp_assign::$method_assign(&mut self, rhs);
81 self
82 }
83 }
84 impl<T, C> $imp<T> for &FormalPowerSeries<T, C>
85 where
86 T: FormalPowerSeriesCoefficient,
87 {
88 type Output = FormalPowerSeries<T, C>;
89 fn $method(self, rhs: T) -> Self::Output {
90 $imp::$method(self.clone(), rhs)
91 }
92 }
93 impl<T, C> $imp<&T> for &FormalPowerSeries<T, C>
94 where
95 T: FormalPowerSeriesCoefficient,
96 {
97 type Output = FormalPowerSeries<T, C>;
98 fn $method(self, rhs: &T) -> Self::Output {
99 $imp::$method(self.clone(), rhs)
100 }
101 }
102 };
103}
104impl_fps_single_binop!(Add, add, AddAssign, add_assign);
105impl_fps_single_binop!(Sub, sub, SubAssign, sub_assign);
106impl_fps_single_binop!(Mul, mul, MulAssign, mul_assign);
107impl_fps_single_binop!(Div, div, DivAssign, div_assign);
108
109impl<T, C> AddAssign<&Self> for FormalPowerSeries<T, C>
110where
111 T: FormalPowerSeriesCoefficient,
112{
113 fn add_assign(&mut self, rhs: &Self) {
114 if self.length() < rhs.length() {
115 self.resize(rhs.length());
116 }
117 for (x, y) in self.iter_mut().zip(rhs.iter()) {
118 x.add_assign(y);
119 }
120 }
121}
122impl<T, C> SubAssign<&Self> for FormalPowerSeries<T, C>
123where
124 T: FormalPowerSeriesCoefficient,
125{
126 fn sub_assign(&mut self, rhs: &Self) {
127 if self.length() < rhs.length() {
128 self.resize(rhs.length());
129 }
130 for (x, y) in self.iter_mut().zip(rhs.iter()) {
131 x.sub_assign(y);
132 }
133 self.trim_tail_zeros();
134 }
135}
136
137macro_rules! impl_fps_binop_addsub {
138 ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
139 impl<T, C> $imp_assign for FormalPowerSeries<T, C>
140 where
141 T: FormalPowerSeriesCoefficient,
142 {
143 fn $method_assign(&mut self, rhs: Self) {
144 $imp_assign::$method_assign(self, &rhs);
145 }
146 }
147 impl<T, C> $imp for FormalPowerSeries<T, C>
148 where
149 T: FormalPowerSeriesCoefficient,
150 {
151 type Output = Self;
152 fn $method(mut self, rhs: Self) -> Self::Output {
153 $imp_assign::$method_assign(&mut self, &rhs);
154 self
155 }
156 }
157 impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
158 where
159 T: FormalPowerSeriesCoefficient,
160 {
161 type Output = Self;
162 fn $method(mut self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
163 $imp_assign::$method_assign(&mut self, rhs);
164 self
165 }
166 }
167 impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
168 where
169 T: FormalPowerSeriesCoefficient,
170 {
171 type Output = FormalPowerSeries<T, C>;
172 fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
173 let mut self_ = self.clone();
174 $imp_assign::$method_assign(&mut self_, &rhs);
175 self_
176 }
177 }
178 impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
179 where
180 T: FormalPowerSeriesCoefficient,
181 {
182 type Output = FormalPowerSeries<T, C>;
183 fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
184 let mut self_ = self.clone();
185 $imp_assign::$method_assign(&mut self_, rhs);
186 self_
187 }
188 }
189 };
190}
191impl_fps_binop_addsub!(Add, add, AddAssign, add_assign);
192impl_fps_binop_addsub!(Sub, sub, SubAssign, sub_assign);
193
194impl<T, C> Mul for FormalPowerSeries<T, C>
195where
196 C: ConvolveSteps<T = Vec<T>>,
197{
198 type Output = Self;
199 fn mul(self, rhs: Self) -> Self::Output {
200 Self::from_vec(C::convolve(self.data, rhs.data))
201 }
202}
203impl<T, C> Div for FormalPowerSeries<T, C>
204where
205 T: FormalPowerSeriesCoefficient,
206 C: ConvolveSteps<T = Vec<T>>,
207{
208 type Output = Self;
209 fn div(mut self, mut rhs: Self) -> Self::Output {
210 self.trim_tail_zeros();
211 rhs.trim_tail_zeros();
212 if self.length() < rhs.length() {
213 return Self::zero();
214 }
215 self.data.reverse();
216 rhs.data.reverse();
217 let n = self.length() - rhs.length() + 1;
218 let mut res = self * rhs.inv(n);
219 res.truncate(n);
220 res.data.reverse();
221 res
222 }
223}
224impl<T, C> Rem for FormalPowerSeries<T, C>
225where
226 T: FormalPowerSeriesCoefficient,
227 C: ConvolveSteps<T = Vec<T>>,
228{
229 type Output = Self;
230 fn rem(self, rhs: Self) -> Self::Output {
231 let mut rem = self.clone() - self / rhs.clone() * rhs;
232 rem.trim_tail_zeros();
233 rem
234 }
235}
236
237impl<T, C> FormalPowerSeries<T, C>
238where
239 T: FormalPowerSeriesCoefficient,
240 C: ConvolveSteps<T = Vec<T>>,
241{
242 pub fn div_rem(self, rhs: Self) -> (Self, Self) {
243 let div = self.clone() / rhs.clone();
244 let mut rem = self - div.clone() * rhs;
245 rem.trim_tail_zeros();
246 (div, rem)
247 }crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 330)
304fn brute_force<T, C>(
305 mut p: FormalPowerSeries<T, C>,
306 mut q: FormalPowerSeries<T, C>,
307 k: usize,
308) -> FpsMatrix<T, C>
309where
310 T: FormalPowerSeriesCoefficient,
311 C: NttReuse<T = Vec<T>>,
312 C::F: Clone,
313{
314 let threshold = degree(&p) - k as isize;
315 let mut matrix = FpsMatrix::identity();
316 while q.length() as isize > threshold {
317 let q_degree = q.length() - 1;
318 let mut negative_quotient = vec![T::zero(); p.length() - q.length() + 1];
319 let inverse = -T::one() / &q[q_degree];
320 for i in (0..negative_quotient.len()).rev() {
321 negative_quotient[i] = p[i + q_degree].clone() * &inverse;
322 p[i + q_degree] = T::zero();
323 for j in 0..q_degree {
324 let value = negative_quotient[i].clone() * &q[j];
325 p[i + j] += &value;
326 }
327 }
328 matrix.left_multiply_step(&negative_quotient);
329 p.truncate(q_degree);
330 p.trim_tail_zeros();
331 swap(&mut p, &mut q);
332 }
333 matrix
334}
335
336fn reduced_transform<T, C>(fps: &FormalPowerSeries<T, C>, length: usize) -> C::F
337where
338 T: FormalPowerSeriesCoefficient,
339 C: NttReuse<T = Vec<T>>,
340{
341 let mut coefficients = vec![T::zero(); length];
342 for (i, value) in fps.iter().enumerate() {
343 coefficients[i & (length - 1)] += value;
344 }
345 C::transform_ntt(coefficients, length)
346}
347
348fn transform_window<T, C>(fps: &FormalPowerSeries<T, C>, end: isize, length: usize) -> C::F
349where
350 T: FormalPowerSeriesCoefficient,
351 C: NttReuse<T = Vec<T>>,
352{
353 let start = end - length as isize;
354 let coefficients = (start..end).map(|index| coefficient(fps, index)).collect();
355 C::transform_ntt(coefficients, length)
356}
357
358fn half_gcd<T, C>(
359 p: &FormalPowerSeries<T, C>,
360 q: &FormalPowerSeries<T, C>,
361 k: usize,
362 length: usize,
363) -> (FpsMatrix<T, C>, FrequencyMatrix<C>)
364where
365 T: FormalPowerSeriesCoefficient,
366 C: NttReuse<T = Vec<T>>,
367 C::F: Clone,
368{
369 let d = degree(p);
370 if degree(q) < d - k as isize {
371 let matrix = FpsMatrix::identity();
372 let frequency = matrix.transform(length);
373 return (matrix, frequency);
374 }
375 if k == 1 {
376 let matrix = FpsMatrix {
377 a00: FormalPowerSeries::zero(),
378 a01: FormalPowerSeries::one(),
379 a10: FormalPowerSeries::one(),
380 a11: -(tail(p, d - 2) / tail(q, d - 2)),
381 };
382 let frequency = matrix.transform(length);
383 return (matrix, frequency);
384 }
385 if p.length().min(q.length()) <= 32 {
386 let matrix = brute_force(p.clone(), q.clone(), k);
387 let frequency = matrix.transform(length);
388 return (matrix, frequency);
389 }
390
391 let half = length / 2;
392 if k <= half {
393 let (matrix, frequency) = half_gcd(p, q, k, half);
394 let frequency = matrix.extend_transform(frequency, length);
395 return (matrix, frequency);
396 }
397
398 let (matrix, mut matrix_frequency) = half_gcd(
399 &tail(p, d - 2 * half as isize),
400 &tail(q, d - 2 * half as isize),
401 half,
402 length,
403 );
404 let degeneracy = half as isize - degree(&matrix.a11);
405
406 let (p0, q0) = matrix_frequency.apply(
407 &transform_window(p, d - half as isize + degeneracy, length),
408 &transform_window(q, d - half as isize + degeneracy, length),
409 length,
410 );
411 let (p1, q1) = matrix_frequency.apply(
412 &transform_window(p, d - 2 * half as isize, length),
413 &transform_window(q, d - 2 * half as isize, length),
414 length,
415 );
416 let part_length = (half as isize + degeneracy) as usize;
417 let mut p_reduced = p1[length - part_length..].to_vec();
418 p_reduced.extend_from_slice(&p0[length - part_length..]);
419 let mut q_reduced = q1[length - part_length..].to_vec();
420 q_reduced.extend_from_slice(&q0[length - part_length..]);
421 let mut q_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(q_reduced).trimed();
422
423 let position = d - half as isize + degeneracy;
424 let mut leading = T::zero();
425 for i in 0..=position {
426 leading += coefficient(p, i) * coefficient(&matrix.a00, position - i)
427 + coefficient(q, i) * coefficient(&matrix.a01, position - i);
428 }
429 p_reduced.push(leading);
430 let mut p_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(p_reduced);
431 if degree(&q_reduced) < 3 * half as isize + degeneracy - k as isize {
432 return (matrix, matrix_frequency);
433 }
434
435 let mut remaining = k as isize - degree(&matrix.a11);
436 let mut top_product = matrix.a11.data.last().unwrap().clone();
437 let mut product_degree = degree(&matrix.a11);
438 if degeneracy > 0 {
439 let skip = (2 * half as isize + 2 * degeneracy - (d - half as isize + degeneracy)).max(0);
440 let (division, remainder) = tail(&p_reduced, skip).div_rem(tail(&q_reduced, skip));
441 remaining -= degree(&division);
442 top_product *= -division.data.last().unwrap().clone();
443 product_degree += degree(&division);
444 matrix_frequency = matrix_frequency.left_multiply_step(&division, length);
445 swap(&mut p_reduced, &mut q_reduced);
446 q_reduced = FormalPowerSeries::zeros(skip as usize);
447 q_reduced.data.extend(remainder.data);
448 }
449
450 let start = 3 * half as isize + degeneracy - k as isize - remaining;
451 let (right_matrix, right_frequency) = half_gcd(
452 &tail(&p_reduced, start),
453 &tail(&q_reduced, start),
454 remaining as usize,
455 length,
456 );
457 let product_frequency = right_frequency.multiply(&matrix_frequency);
458 let mut product = product_frequency.clone().inverse_transform(length);
459 product.a00.truncate(k);
460 product.a00.trim_tail_zeros();
461 product.a01.truncate(k);
462 product.a01.trim_tail_zeros();
463 product.a10.truncate(k);
464 product.a10.trim_tail_zeros();
465 product_degree += degree(&right_matrix.a11);
466 if product_degree == length as isize {
467 product.a11.resize(k + 1);
468 let highest = top_product * right_matrix.a11.data.last().unwrap();
469 product.a11[k] = highest.clone();
470 product.a11[0] -= highest;
471 }
472 product.a11.trim_tail_zeros();
473 let product_frequency = if C::MULTIPLE {
474 product.transform(length)
475 } else {
476 product_frequency
477 };
478 (product, product_frequency)
479}Sourcepub fn trimed(self) -> Self
pub fn trimed(self) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 237)
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 }
265}
266
267fn degree<T, C>(fps: &FormalPowerSeries<T, C>) -> isize {
268 fps.length() as isize - 1
269}
270
271fn add<T, C>(
272 left: FormalPowerSeries<T, C>,
273 right: FormalPowerSeries<T, C>,
274) -> FormalPowerSeries<T, C>
275where
276 T: FormalPowerSeriesCoefficient,
277{
278 (left + right).trimed()
279}
280
281fn tail<T, C>(fps: &FormalPowerSeries<T, C>, start: isize) -> FormalPowerSeries<T, C>
282where
283 T: FormalPowerSeriesCoefficient,
284{
285 let start = start.max(0) as usize;
286 if start >= fps.length() {
287 FormalPowerSeries::zero()
288 } else {
289 FormalPowerSeries::from_vec(fps.data[start..].to_vec())
290 }
291}
292
293fn coefficient<T, C>(fps: &FormalPowerSeries<T, C>, index: isize) -> T
294where
295 T: FormalPowerSeriesCoefficient,
296{
297 if index < 0 {
298 T::zero()
299 } else {
300 fps.coeff(index as usize)
301 }
302}
303
304fn brute_force<T, C>(
305 mut p: FormalPowerSeries<T, C>,
306 mut q: FormalPowerSeries<T, C>,
307 k: usize,
308) -> FpsMatrix<T, C>
309where
310 T: FormalPowerSeriesCoefficient,
311 C: NttReuse<T = Vec<T>>,
312 C::F: Clone,
313{
314 let threshold = degree(&p) - k as isize;
315 let mut matrix = FpsMatrix::identity();
316 while q.length() as isize > threshold {
317 let q_degree = q.length() - 1;
318 let mut negative_quotient = vec![T::zero(); p.length() - q.length() + 1];
319 let inverse = -T::one() / &q[q_degree];
320 for i in (0..negative_quotient.len()).rev() {
321 negative_quotient[i] = p[i + q_degree].clone() * &inverse;
322 p[i + q_degree] = T::zero();
323 for j in 0..q_degree {
324 let value = negative_quotient[i].clone() * &q[j];
325 p[i + j] += &value;
326 }
327 }
328 matrix.left_multiply_step(&negative_quotient);
329 p.truncate(q_degree);
330 p.trim_tail_zeros();
331 swap(&mut p, &mut q);
332 }
333 matrix
334}
335
336fn reduced_transform<T, C>(fps: &FormalPowerSeries<T, C>, length: usize) -> C::F
337where
338 T: FormalPowerSeriesCoefficient,
339 C: NttReuse<T = Vec<T>>,
340{
341 let mut coefficients = vec![T::zero(); length];
342 for (i, value) in fps.iter().enumerate() {
343 coefficients[i & (length - 1)] += value;
344 }
345 C::transform_ntt(coefficients, length)
346}
347
348fn transform_window<T, C>(fps: &FormalPowerSeries<T, C>, end: isize, length: usize) -> C::F
349where
350 T: FormalPowerSeriesCoefficient,
351 C: NttReuse<T = Vec<T>>,
352{
353 let start = end - length as isize;
354 let coefficients = (start..end).map(|index| coefficient(fps, index)).collect();
355 C::transform_ntt(coefficients, length)
356}
357
358fn half_gcd<T, C>(
359 p: &FormalPowerSeries<T, C>,
360 q: &FormalPowerSeries<T, C>,
361 k: usize,
362 length: usize,
363) -> (FpsMatrix<T, C>, FrequencyMatrix<C>)
364where
365 T: FormalPowerSeriesCoefficient,
366 C: NttReuse<T = Vec<T>>,
367 C::F: Clone,
368{
369 let d = degree(p);
370 if degree(q) < d - k as isize {
371 let matrix = FpsMatrix::identity();
372 let frequency = matrix.transform(length);
373 return (matrix, frequency);
374 }
375 if k == 1 {
376 let matrix = FpsMatrix {
377 a00: FormalPowerSeries::zero(),
378 a01: FormalPowerSeries::one(),
379 a10: FormalPowerSeries::one(),
380 a11: -(tail(p, d - 2) / tail(q, d - 2)),
381 };
382 let frequency = matrix.transform(length);
383 return (matrix, frequency);
384 }
385 if p.length().min(q.length()) <= 32 {
386 let matrix = brute_force(p.clone(), q.clone(), k);
387 let frequency = matrix.transform(length);
388 return (matrix, frequency);
389 }
390
391 let half = length / 2;
392 if k <= half {
393 let (matrix, frequency) = half_gcd(p, q, k, half);
394 let frequency = matrix.extend_transform(frequency, length);
395 return (matrix, frequency);
396 }
397
398 let (matrix, mut matrix_frequency) = half_gcd(
399 &tail(p, d - 2 * half as isize),
400 &tail(q, d - 2 * half as isize),
401 half,
402 length,
403 );
404 let degeneracy = half as isize - degree(&matrix.a11);
405
406 let (p0, q0) = matrix_frequency.apply(
407 &transform_window(p, d - half as isize + degeneracy, length),
408 &transform_window(q, d - half as isize + degeneracy, length),
409 length,
410 );
411 let (p1, q1) = matrix_frequency.apply(
412 &transform_window(p, d - 2 * half as isize, length),
413 &transform_window(q, d - 2 * half as isize, length),
414 length,
415 );
416 let part_length = (half as isize + degeneracy) as usize;
417 let mut p_reduced = p1[length - part_length..].to_vec();
418 p_reduced.extend_from_slice(&p0[length - part_length..]);
419 let mut q_reduced = q1[length - part_length..].to_vec();
420 q_reduced.extend_from_slice(&q0[length - part_length..]);
421 let mut q_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(q_reduced).trimed();
422
423 let position = d - half as isize + degeneracy;
424 let mut leading = T::zero();
425 for i in 0..=position {
426 leading += coefficient(p, i) * coefficient(&matrix.a00, position - i)
427 + coefficient(q, i) * coefficient(&matrix.a01, position - i);
428 }
429 p_reduced.push(leading);
430 let mut p_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(p_reduced);
431 if degree(&q_reduced) < 3 * half as isize + degeneracy - k as isize {
432 return (matrix, matrix_frequency);
433 }
434
435 let mut remaining = k as isize - degree(&matrix.a11);
436 let mut top_product = matrix.a11.data.last().unwrap().clone();
437 let mut product_degree = degree(&matrix.a11);
438 if degeneracy > 0 {
439 let skip = (2 * half as isize + 2 * degeneracy - (d - half as isize + degeneracy)).max(0);
440 let (division, remainder) = tail(&p_reduced, skip).div_rem(tail(&q_reduced, skip));
441 remaining -= degree(&division);
442 top_product *= -division.data.last().unwrap().clone();
443 product_degree += degree(&division);
444 matrix_frequency = matrix_frequency.left_multiply_step(&division, length);
445 swap(&mut p_reduced, &mut q_reduced);
446 q_reduced = FormalPowerSeries::zeros(skip as usize);
447 q_reduced.data.extend(remainder.data);
448 }
449
450 let start = 3 * half as isize + degeneracy - k as isize - remaining;
451 let (right_matrix, right_frequency) = half_gcd(
452 &tail(&p_reduced, start),
453 &tail(&q_reduced, start),
454 remaining as usize,
455 length,
456 );
457 let product_frequency = right_frequency.multiply(&matrix_frequency);
458 let mut product = product_frequency.clone().inverse_transform(length);
459 product.a00.truncate(k);
460 product.a00.trim_tail_zeros();
461 product.a01.truncate(k);
462 product.a01.trim_tail_zeros();
463 product.a10.truncate(k);
464 product.a10.trim_tail_zeros();
465 product_degree += degree(&right_matrix.a11);
466 if product_degree == length as isize {
467 product.a11.resize(k + 1);
468 let highest = top_product * right_matrix.a11.data.last().unwrap();
469 product.a11[k] = highest.clone();
470 product.a11[0] -= highest;
471 }
472 product.a11.trim_tail_zeros();
473 let product_frequency = if C::MULTIPLE {
474 product.transform(length)
475 } else {
476 product_frequency
477 };
478 (product, product_frequency)
479}Source§impl<T, C> FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Sourcepub fn prefix_ref(&self, deg: usize) -> Self
pub fn prefix_ref(&self, deg: usize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 345)
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }Sourcepub fn prefix(self, deg: usize) -> Self
pub fn prefix(self, deg: usize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 610)
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }Sourcepub fn even(self) -> Self
pub fn even(self) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 1064)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }pub fn odd(self) -> Self
Sourcepub fn diff(self) -> Self
pub fn diff(self) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 345)
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }Sourcepub fn integral(self) -> Self
pub fn integral(self) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 605)
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }Sourcepub fn parity_inversion(self) -> Self
pub fn parity_inversion(self) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 1063)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }Sourcepub fn eval(&self, x: T) -> T
pub fn eval(&self, x: T) -> T
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 1095)
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }Source§impl<T, C> FormalPowerSeries<T, C>
impl<T, C> FormalPowerSeries<T, C>
Sourcefn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize>
fn sparse_stride(&self, deg: usize, factor: usize) -> Option<usize>
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 271)
266 pub fn inv(&self, deg: usize) -> Self {
267 if deg == 0 {
268 return Self::zero();
269 }
270 debug_assert!(!self[0].is_zero());
271 if let Some(step) = self.sparse_stride(deg, 6) {
272 let inv = T::one() / self[0].clone();
273 let pos: Vec<_> = self
274 .data
275 .iter()
276 .take(deg)
277 .enumerate()
278 .skip(1)
279 .filter(|(_, x)| !x.is_zero())
280 .map(|(i, x)| (i, -x.clone() * &inv))
281 .collect();
282 let mut f = Self::zeros(deg);
283 f[0] = inv;
284 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
285 let mut tot = T::zero();
286 for (j, coefficient) in &pos {
287 if *j > i {
288 break;
289 }
290 tot += coefficient.clone() * &f[i - *j];
291 }
292 f[i] = tot;
293 }
294 return f;
295 }
296 let mut f = Self::from(T::one() / self[0].clone());
297 f.data.reserve(deg.saturating_sub(1));
298 let extend = |f: &mut Self, end| {
299 for i in f.length()..end {
300 let mut tot = T::zero();
301 for j in 1..=i.min(self.length() - 1) {
302 tot += self[j].clone() * &f[i - j];
303 }
304 f.data.push(-tot * &f[0]);
305 }
306 };
307 extend(&mut f, deg.min(32));
308 let mut error = Vec::new();
309 let mut i = f.length();
310 while i < deg {
311 if deg - i <= 4 {
312 extend(&mut f, deg);
313 break;
314 }
315 error.clear();
316 error.extend(
317 self.data[..(i * 2).min(deg).min(self.length())]
318 .iter()
319 .cloned(),
320 );
321 let factor = C::transform(f.data.clone(), 2 * i);
322 let mut error_fft = C::transform(error, 2 * i);
323 C::multiply(&mut error_fft, &factor);
324 error = C::inverse_transform(error_fft, 2 * i);
325 error.drain(..i);
326 let mut error_fft = C::transform(error, 2 * i);
327 C::multiply(&mut error_fft, &factor);
328 error = C::inverse_transform(error_fft, 2 * i);
329 error.truncate(i.min(deg - i));
330 f.data.extend(error.drain(..).map(Neg::neg));
331 i *= 2;
332 }
333 f
334 }
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }Sourcepub fn inv(&self, deg: usize) -> Self
pub fn inv(&self, deg: usize) -> Self
Examples found in repository?
More examples
crates/library_checker/src/polynomial/inv_of_formal_power_series_sparse.rs (line 14)
5pub fn inv_of_formal_power_series_sparse(reader: impl Read, writer: impl Write) {
6 prepare_io!(reader, writer);
7 sc!(n, k);
8 let mut a = vec![M::zero(); n];
9 for _ in 0..k {
10 sc!(i, a_i: M);
11 a[i] = a_i;
12 }
13 let f = Fps998244353::from_vec(a);
14 let g = f.inv(n);
15 pp!(@it g.data);
16}crates/competitive/src/math/formal_power_series/formal_power_series_nums.rs (line 218)
209 fn div(mut self, mut rhs: Self) -> Self::Output {
210 self.trim_tail_zeros();
211 rhs.trim_tail_zeros();
212 if self.length() < rhs.length() {
213 return Self::zero();
214 }
215 self.data.reverse();
216 rhs.data.reverse();
217 let n = self.length() - rhs.length() + 1;
218 let mut res = self * rhs.inv(n);
219 res.truncate(n);
220 res.data.reverse();
221 res
222 }crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 422)
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }Sourcepub fn exp(&self, deg: usize) -> Self
pub fn exp(&self, deg: usize) -> Self
Examples found in repository?
More examples
crates/library_checker/src/polynomial/exp_of_formal_power_series_sparse.rs (line 14)
5pub fn exp_of_formal_power_series_sparse(reader: impl Read, writer: impl Write) {
6 prepare_io!(reader, writer);
7 sc!(n, k);
8 let mut a = vec![M::zero(); n];
9 for _ in 0..k {
10 sc!(i, a_i: M);
11 a[i] = a_i;
12 }
13 let f = Fps998244353::from_vec(a);
14 let g = f.exp(n);
15 pp!(@it g.data);
16}crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 992)
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }Sourcefn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 472)
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }Sourcefn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 370)
335 pub fn exp(&self, deg: usize) -> Self
336 where
337 C: NttReuse<T = Vec<T>>,
338 C::F: Clone,
339 {
340 if deg == 0 {
341 return Self::zero();
342 }
343 debug_assert!(self[0].is_zero());
344 if let Some(step) = self.sparse_stride(deg, if deg <= 256 { 16 } else { 8 }) {
345 let diff = self.prefix_ref(deg).diff();
346 let pos: Vec<_> = diff
347 .data
348 .iter()
349 .enumerate()
350 .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
351 .collect();
352 let mut f = Self::zeros(deg);
353 f[0] = T::one();
354 if pos.is_empty() {
355 return f;
356 }
357 let mf = T::memorized_factorial(deg);
358 for i in (pos.first().map_or(deg, |j| j + 1)..deg).step_by(step) {
359 let mut tot = T::zero();
360 for &j in &pos {
361 if j > i - 1 {
362 break;
363 }
364 tot += f[i - 1 - j].clone() * &diff[j];
365 }
366 f[i] = tot * T::memorized_inv(&mf, i);
367 }
368 return f;
369 }
370 self.exp_or_pow(None, deg)
371 }
372
373 fn sum_products(f: &[C::F], g: &[C::F], len: usize) -> Vec<T>
374 where
375 C: NttReuse<T = Vec<T>>,
376 C::F: Clone,
377 {
378 let chunk = C::max_product_sum_count(&f[0]);
379 f.rchunks(chunk)
380 .zip(g.chunks(chunk))
381 .map(|(f, g)| {
382 let mut sum = f[f.len() - 1].clone();
383 C::multiply_prefix(&mut sum, &g[0]);
384 for (f, g) in f.iter().rev().skip(1).zip(&g[1..]) {
385 C::multiply_add(&mut sum, f, g);
386 }
387 C::inverse_transform_ntt(sum, len)
388 })
389 .reduce(|mut sum, part| {
390 for (sum, value) in sum.iter_mut().zip(part) {
391 *sum += value;
392 }
393 sum
394 })
395 .unwrap()
396 }
397
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }Sourcefn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 420)
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }Sourcepub fn log(&self, deg: usize) -> Self
pub fn log(&self, deg: usize) -> Self
Examples found in repository?
More examples
crates/library_checker/src/polynomial/log_of_formal_power_series_sparse.rs (line 14)
5pub fn log_of_formal_power_series_sparse(reader: impl Read, writer: impl Write) {
6 prepare_io!(reader, writer);
7 sc!(n, k);
8 let mut a = vec![M::zero(); n];
9 for _ in 0..k {
10 sc!(i, a_i: M);
11 a[i] = a_i;
12 }
13 let f = Fps998244353::from_vec(a);
14 let g = f.log(n);
15 pp!(@it g.data);
16}crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 415)
398 fn exp_or_pow(&self, power: Option<T>, deg: usize) -> Self
399 where
400 C: NttReuse<T = Vec<T>>,
401 C::F: Clone,
402 {
403 if deg == 1 {
404 return Self::one();
405 }
406 let indices: Vec<_> = (0..=deg).map(T::from).collect();
407 let modulus = <T::Base as MIntConvert<usize>>::mod_into();
408 let mut inv = vec![T::zero(); deg + 1];
409 inv[1] = T::one();
410 for i in 2..=deg {
411 inv[i] = -T::from(modulus / i) * &inv[modulus % i];
412 }
413 let block = deg.next_power_of_two() / 16;
414 let logarithm = if let Some(rhs) = &power {
415 self.prefix_ref(block).log(block) * rhs
416 } else {
417 self.prefix_ref(block)
418 };
419 let (kernel, mut kernel_inverse, previous_inverse_fft) =
420 logarithm.exp_newton(block, &indices, &inv);
421 if power.is_some() {
422 kernel_inverse = (self.prefix_ref(block) * &kernel).inv(block);
423 } else {
424 let mut error_fft = C::transform_ntt(kernel.data.clone(), block);
425 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
426 let error = C::inverse_transform_ntt(error_fft, block);
427 let mut error_fft =
428 C::transform_ntt(error.into_iter().skip(block / 2).collect(), block);
429 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
430 let error = C::inverse_transform_ntt(error_fft, block / 2);
431 kernel_inverse
432 .data
433 .extend(error.into_iter().take(block / 2).map(Neg::neg));
434 }
435 let kernel_data = kernel.data;
436 let kernel_inverse_data = kernel_inverse.data;
437 let kernel_inverse = C::transform(kernel_inverse_data, block * 2);
438 let kernel = C::transform(kernel_data.clone(), block * 2);
439 let blocks = deg.div_ceil(block);
440 let mut derivative_ffts = Vec::with_capacity(blocks - 1);
441 let mut polynomial_ffts = Vec::with_capacity(if power.is_some() { blocks - 1 } else { 0 });
442 for q in 1..blocks {
443 let mut values = Self::zeros(block * 2);
444 for (i, values) in values.data.chunks_mut(block).enumerate() {
445 let start = (q - i) * block;
446 for (value, x) in values
447 .iter_mut()
448 .zip(self.iter().skip(start).take(deg - start))
449 {
450 *value = x.clone();
451 }
452 }
453 if power.is_some() {
454 polynomial_ffts.push(C::transform_ntt(values.data.clone(), block * 2));
455 }
456 for (i, values) in values.data.chunks_mut(block).enumerate() {
457 let start = (q - i) * block;
458 for (value, index) in values.iter_mut().zip(&indices[start..]) {
459 *value *= index;
460 }
461 }
462 derivative_ffts.push(C::transform_ntt(values.data, block * 2));
463 }
464 let mut result = kernel_data.clone();
465 result.reserve(deg - block);
466 let mut result_ffts = Vec::with_capacity(blocks - 1);
467 for q in 1..blocks {
468 result_ffts.push(C::transform_ntt(
469 result[(q - 1) * block..q * block].to_vec(),
470 block * 2,
471 ));
472 let mut values = Self::sum_products(&derivative_ffts[..q], &result_ffts, block);
473 if let Some(rhs) = &power {
474 let product = Self::sum_products(&polynomial_ffts[..q], &result_ffts, block);
475 let factor = rhs.clone() + T::one();
476 // The power satisfies f g' = rhs f' g.
477 for (i, value) in values.iter_mut().take(deg - q * block).enumerate() {
478 *value = value.clone() * &factor - product[i].clone() * &indices[q * block + i];
479 }
480 }
481 let mut values = C::transform(values, block * 2);
482 C::multiply(&mut values, &kernel_inverse);
483 let mut values = C::inverse_transform(values, block * 2);
484 values.truncate(block);
485 let len = block.min(deg - q * block);
486 for (i, value) in values.iter_mut().take(len).enumerate() {
487 *value *= &inv[q * block + i];
488 }
489 values[len..].fill(T::zero());
490 let mut values = C::transform(values, block * 2);
491 C::multiply(&mut values, &kernel);
492 let mut values = C::inverse_transform(values, block * 2);
493 values.truncate(len);
494 result.extend(values);
495 }
496 Self::from_vec(result)
497 }
498
499 fn exp_newton(&self, deg: usize, indices: &[T], inv: &[T]) -> (Self, Self, C::F)
500 where
501 C: NttReuse<T = Vec<T>>,
502 C::F: Clone,
503 {
504 if deg == 1 {
505 let one = Self::one();
506 return (one.clone(), one.clone(), C::transform_ntt(one.data, 1));
507 }
508 let mut f = Self::from_vec(vec![T::one(), self.coeff(1)]);
509 let mut inverse = Self::one();
510 let mut inverse_fft = C::transform_ntt(inverse.data.clone(), 2);
511 let mut m = 2;
512 while m < deg {
513 let f_fft = C::transform_ntt(f.data.clone(), 2 * m);
514
515 let previous_inverse_fft = inverse_fft;
516 let mut error_fft = previous_inverse_fft.clone();
517 C::multiply_prefix(&mut error_fft, &f_fft);
518 let mut error = C::inverse_transform_ntt(error_fft, m);
519 error[..m / 2].fill(T::zero());
520 let mut error_fft = C::transform_ntt(error, m);
521 C::multiply_prefix(&mut error_fft, &previous_inverse_fft);
522 let error = C::inverse_transform_ntt(error_fft, m);
523 inverse
524 .data
525 .extend(error.into_iter().skip(m / 2).map(Neg::neg));
526 inverse_fft = C::transform_ntt(inverse.data.clone(), 2 * m);
527
528 let mut delta = Self::from_vec(
529 self.data
530 .iter()
531 .take(m)
532 .enumerate()
533 .skip(1)
534 .map(|(i, value)| value.clone() * &indices[i])
535 .collect(),
536 );
537 delta.resize(m);
538 let mut delta_fft = C::transform_ntt(delta.data, m);
539 C::multiply_prefix(&mut delta_fft, &f_fft);
540 let mut delta = Self::from_vec(C::inverse_transform_ntt(delta_fft, m));
541 for i in 1..f.length() {
542 delta[i - 1] -= f[i].clone() * &indices[i];
543 }
544 delta.resize(2 * m);
545 for i in (0..m - 1).rev() {
546 delta.data[m + i] = delta.data[i].clone();
547 }
548 delta.data[..m - 1].fill(T::zero());
549 let mut delta_fft = C::transform_ntt(delta.data, 2 * m);
550 C::multiply_prefix(&mut delta_fft, &inverse_fft);
551 let mut delta = C::inverse_transform_ntt(delta_fft, 2 * m);
552 delta.pop();
553 delta.push(T::zero());
554 let target = (2 * m).min(deg);
555 for i in (1..target).rev() {
556 delta[i] = delta[i - 1].clone() * &inv[i];
557 }
558 delta[0] = T::zero();
559 delta[target..].fill(T::zero());
560 for i in m..(2 * m).min(self.length()) {
561 delta[i] += self[i].clone();
562 }
563 delta[..m].fill(T::zero());
564 let mut delta_fft = C::transform_ntt(delta, 2 * m);
565 C::multiply_prefix(&mut delta_fft, &f_fft);
566 let delta = C::inverse_transform_ntt(delta_fft, 2 * m);
567 f.data
568 .extend(delta.into_iter().skip(m).take((deg - m).min(m)));
569 m *= 2;
570 }
571 (f, inverse, inverse_fft)
572 }
573 pub fn log(&self, deg: usize) -> Self {
574 if deg == 0 {
575 return Self::zero();
576 }
577 debug_assert!(!self[0].is_zero());
578 if deg == 1 {
579 return Self::zeros(1);
580 }
581 if let Some(step) = self.sparse_stride(deg, 2) {
582 let pos: Vec<_> = self
583 .iter()
584 .take(deg)
585 .enumerate()
586 .skip(1)
587 .filter_map(|(i, x)| (!x.is_zero()).then_some(i))
588 .collect();
589 let mut derivative = Self::zeros(deg);
590 let inverse = T::one() / self[0].clone();
591 for i in (pos.first().copied().unwrap_or(deg)..deg).step_by(step) {
592 let mut value = self.coeff(i) * T::from(i);
593 for &j in &pos {
594 if j >= i {
595 break;
596 }
597 value -= self[j].clone() * &derivative[i - j];
598 }
599 derivative[i] = value * &inverse;
600 }
601 if pos.is_empty() {
602 return derivative;
603 }
604 derivative.data.remove(0);
605 return derivative.integral();
606 }
607 let n = deg - 1;
608 if n <= 64 {
609 return (self.inv(deg) * self.prefix_ref(deg).diff())
610 .prefix(n)
611 .integral();
612 }
613 let half = n.next_power_of_two() / 2;
614 let derivative = self.prefix_ref(deg).diff();
615 let inverse = C::transform(self.inv(half).data, half * 2);
616 let mut quotient = C::transform(derivative.prefix_ref(half).data, half * 2);
617 C::multiply(&mut quotient, &inverse);
618 let mut result = C::inverse_transform(quotient, half * 2);
619 result.truncate(half);
620 if n - half <= 4 {
621 let inverse = T::one() / self[0].clone();
622 for i in half..n {
623 let mut value = derivative.coeff(i);
624 for j in 1..=i.min(self.length() - 1) {
625 value -= self[j].clone() * &result[i - j];
626 }
627 result.push(value * &inverse);
628 }
629 return Self::from_vec(result).integral();
630 }
631 let quotient = C::transform(result.clone(), half * 2);
632 let mut error = C::transform(self.prefix_ref(n).data, half * 2);
633 C::multiply(&mut error, "ient);
634 let mut error = C::inverse_transform(error, half * 2);
635 for i in 0..n - half {
636 error[i] = derivative.coeff(half + i) - &error[half + i];
637 }
638 error.truncate(n - half);
639 let mut error = C::transform(error, half * 2);
640 C::multiply(&mut error, &inverse);
641 let error = C::inverse_transform(error, half * 2);
642 result.extend(error.into_iter().take(n - half));
643 Self::from_vec(result).integral()
644 }
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }
966}
967
968impl<T, C> FormalPowerSeries<T, C>
969where
970 T: FormalPowerSeriesCoefficient,
971 C: ConvolveSteps<T = Vec<T>>,
972{
973 pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
974 where
975 F: FnMut(usize) -> T,
976 C: NttReuse<T = Vec<T>>,
977 C::F: Clone,
978 {
979 let n = self.length();
980 let mut f = Self::zeros(n);
981 for i in 1..n {
982 if !self[i].is_zero() {
983 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
984 if j & 1 != 0 {
985 f[d] += self[i].clone() * &inverse(j);
986 } else {
987 f[d] -= self[i].clone() * &inverse(j);
988 }
989 }
990 }
991 }
992 f.exp(deg)
993 }
994 pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
995 where
996 F: FnMut(usize) -> T,
997 C: NttReuse<T = Vec<T>>,
998 C::F: Clone,
999 {
1000 let n = self.length();
1001 let mut f = Self::zeros(n);
1002 for i in 1..n {
1003 if !self[i].is_zero() {
1004 for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
1005 f[d] += self[i].clone() * &inverse(j);
1006 }
1007 }
1008 }
1009 f.exp(deg)
1010 }
1011 /// [x^n] P(x) / Q(x)
1012 pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
1013 where
1014 C: NttReuse<T = Vec<T>>,
1015 {
1016 let mut res = T::zero();
1017 rhs.trim_tail_zeros();
1018 if self.length() >= rhs.length() {
1019 let r = &self / &rhs;
1020 if n < r.length() {
1021 res = r[n].clone();
1022 }
1023 self -= r * &rhs;
1024 self.trim_tail_zeros();
1025 }
1026 let mut k = rhs.length().next_power_of_two();
1027 let mut p = C::transform_ntt(self.data, k * 2);
1028 let mut q = C::transform_ntt(rhs.data, k * 2);
1029 while n > 0 {
1030 let t = C::even_mul_normal_neg(&q, &q);
1031 p = if n.is_multiple_of(2) {
1032 C::even_mul_normal_neg(&p, &q)
1033 } else {
1034 C::odd_mul_normal_neg(&p, &q)
1035 };
1036 q = t;
1037 n /= 2;
1038 if n != 0 {
1039 if n < k / 2 {
1040 p = C::transform_ntt(C::inverse_transform_ntt(p, k / 2), k);
1041 q = C::transform_ntt(C::inverse_transform_ntt(q, k / 2), k);
1042 k /= 2;
1043 } else if C::MULTIPLE {
1044 p = C::transform_ntt(C::inverse_transform_ntt(p, k), k * 2);
1045 q = C::transform_ntt(C::inverse_transform_ntt(q, k), k * 2);
1046 } else {
1047 p = C::ntt_doubling(p, false);
1048 q = C::ntt_doubling(q, false);
1049 }
1050 }
1051 }
1052 let p = C::inverse_transform_ntt(p, k);
1053 let q = C::inverse_transform_ntt(q, k);
1054 res + p[0].clone() / q[0].clone()
1055 }
1056 /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }
1165 pub fn product_all<I>(iter: I, deg: usize) -> Self
1166 where
1167 I: IntoIterator<Item = Self>,
1168 {
1169 let mut heap: BinaryHeap<_> = iter
1170 .into_iter()
1171 .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
1172 .collect();
1173 while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
1174 if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
1175 let z = (x * y).prefix(deg);
1176 heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
1177 } else {
1178 return x;
1179 }
1180 }
1181 Self::one()
1182 }
1183 pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
1184 where
1185 I: IntoIterator<Item = (Self, Self)>,
1186 {
1187 let mut heap: BinaryHeap<_> = iter
1188 .into_iter()
1189 .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
1190 .collect();
1191 while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
1192 if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
1193 let zb = (&xb * &yb).prefix(deg);
1194 let za = (xa * yb + ya * xb).prefix(deg);
1195 heap.push(PartialIgnoredOrd(
1196 Reverse(za.length().max(zb.length())),
1197 (za, zb),
1198 ));
1199 } else {
1200 return (xa, xb);
1201 }
1202 }
1203 (Self::zero(), Self::one())
1204 }
1205 pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
1206 where
1207 C: NttReuse<T = Vec<T>>,
1208 {
1209 if let Some(x) = a.get(k) {
1210 return x.clone();
1211 }
1212 let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
1213 p.bostan_mori(self, k)
1214 }
1215 pub fn kth_term(a: Vec<T>, k: usize) -> T
1216 where
1217 C: NttReuse<T = Vec<T>>,
1218 C::F: Clone,
1219 {
1220 if let Some(x) = a.get(k) {
1221 return x.clone();
1222 }
1223 Self::berlekamp_massey(&a).kth_term_of_linearly_recurrence(a, k)
1224 }
1225 /// sum_i a_i exp(b_i x)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }
1242 /// sum_i (a_i x)^j
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }
1257
1258 pub fn power_projection(&self, w: &[T], m: usize) -> Self
1259 where
1260 C: NttReuse<T = Vec<T>>,
1261 {
1262 if w.is_empty() {
1263 return Self::zeros(m);
1264 }
1265 if m <= 1 {
1266 return Self::from_vec(vec![w[0].clone(); m]);
1267 }
1268
1269 let n0 = w.len();
1270 let mut n = n0.next_power_of_two();
1271 let mut f = self.prefix_ref(n);
1272 f.resize(n);
1273
1274 let base = n * 2;
1275 let mut p_flat = vec![T::zero(); base];
1276 for (i, wi) in w.iter().enumerate() {
1277 p_flat[n - 1 - i] = wi.clone();
1278 }
1279 let mut q_flat = vec![T::zero(); base * 2];
1280 q_flat[0] = T::one();
1281 let q_offset = base;
1282 for (i, fi) in f.iter().enumerate() {
1283 q_flat[q_offset + i] = -fi.clone();
1284 }
1285 let mut py = 1usize;
1286 let mut qy = 2usize;
1287
1288 let y_limit = m;
1289 while n > 1 {
1290 let (mut p, mut q) = C::power_projection_step(p_flat, q_flat, n, py, qy);
1291 let new_py = (py + qy - 1).min(y_limit);
1292 let new_qy = (qy + qy - 1).min(y_limit);
1293 p.resize_with(n * new_py, T::zero);
1294 q.resize_with(n * new_qy, T::zero);
1295
1296 let n2 = n / 2;
1297 for row in p.chunks_exact_mut(n) {
1298 row[n2..].fill_with(T::zero);
1299 }
1300 for row in q.chunks_exact_mut(n) {
1301 row[n2..].fill_with(T::zero);
1302 }
1303 p_flat = p;
1304 q_flat = q;
1305 py = new_py;
1306 qy = new_qy;
1307 n = n2;
1308 }
1309
1310 let base = 2;
1311 let mut p_y = Vec::with_capacity(py);
1312 for y in 0..py {
1313 p_y.push(p_flat[base * y].clone());
1314 }
1315 let mut q_y = Vec::with_capacity(qy);
1316 for y in 0..qy {
1317 q_y.push(q_flat[base * y].clone());
1318 }
1319 (Self::from_vec(p_y) * Self::from_vec(q_y).inv(m)).prefix(m)
1320 }
1321
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }Sourcepub fn pow(&self, rhs: usize, deg: usize) -> Self
pub fn pow(&self, rhs: usize, deg: usize) -> Self
Examples found in repository?
More examples
crates/library_checker/src/polynomial/pow_of_formal_power_series_sparse.rs (line 14)
5pub fn pow_of_formal_power_series_sparse(reader: impl Read, writer: impl Write) {
6 prepare_io!(reader, writer);
7 sc!(n, k, m);
8 let mut a = vec![M::zero(); n];
9 for _ in 0..k {
10 sc!(i, a_i: M);
11 a[i] = a_i;
12 }
13 let f = Fps998244353::from_vec(a);
14 let g = f.pow(m, n);
15 pp!(@it g.data);
16}Sourcefn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self
fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 670)
645 pub fn pow(&self, rhs: usize, deg: usize) -> Self
646 where
647 C: NttReuse<T = Vec<T>>,
648 C::F: Clone,
649 {
650 if rhs == 0 {
651 return Self::from_vec(
652 once(T::one())
653 .chain(repeat_with(T::zero))
654 .take(deg)
655 .collect(),
656 );
657 }
658 if rhs == 1 {
659 return self.prefix_ref(deg).resized(deg);
660 }
661 if let Some(k) = self
662 .iter()
663 .take(deg.div_ceil(rhs))
664 .position(|x| !x.is_zero())
665 {
666 let deg = deg - k * rhs;
667 let x0 = self[k].clone();
668 let mut f = (self.prefix_ref(k + deg) >> k) / &x0;
669 if let Some(step) = f.sparse_stride(deg, 12) {
670 f = f.pow_sparse1(T::from(rhs), deg, step);
671 } else if rhs <= 4 {
672 let squared = (&f * &f).prefix(deg);
673 f = match rhs {
674 2 => squared,
675 3 => (squared * f).prefix(deg),
676 _ => (&squared * &squared).prefix(deg),
677 }
678 .resized(deg);
679 } else {
680 f = f.exp_or_pow(Some(T::from(rhs)), deg);
681 }
682 f *= x0.pow(rhs);
683 f <<= k * rhs;
684 f
685 } else {
686 Self::zeros(deg)
687 }
688 }
689 fn pow_sparse1(&self, rhs: T, deg: usize, step: usize) -> Self {
690 debug_assert!(!self[0].is_zero());
691 let mut pos: Vec<_> = self
692 .data
693 .iter()
694 .take(deg)
695 .enumerate()
696 .skip(1)
697 .filter(|(_, x)| !x.is_zero())
698 .map(|(i, x)| (i, x.clone(), T::from(i) * &rhs * x))
699 .collect();
700 let mut f = Self::zeros(deg);
701 f[0] = T::one();
702 if pos.is_empty() {
703 return f;
704 }
705 let mf = T::memorized_factorial(deg);
706 for (_, coefficient, _) in &mut pos {
707 *coefficient *= T::from(step);
708 }
709 for i in (pos.first().map_or(deg, |x| x.0)..deg).step_by(step) {
710 let mut tot = T::zero();
711 for (j, coefficient, weight) in &mut pos {
712 if *j > i {
713 break;
714 }
715 tot += weight.clone() * &f[i - *j];
716 *weight -= &*coefficient;
717 }
718 f[i] = tot * T::memorized_inv(&mf, i);
719 }
720 f
721 }
722
723 fn sparse_fold(&self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize) -> T {
724 sparse
725 .into_iter()
726 .take_while(|&(i, _)| i <= deg)
727 .fold(T::zero(), |sum, (i, x)| sum + x * self.coeff(deg - i))
728 }
729
730 /// solve: $X(QF)'=\alpha P'(QF)+\beta P(Q'F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }
791
792 /// P^exp_p * Q^exp_q
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }
876}
877
878impl<T, C> FormalPowerSeries<T, C>
879where
880 T: FormalPowerSeriesCoefficientSqrt,
881 C: ConvolveSteps<T = Vec<T>>,
882{
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }Sourcefn sparse_fold(
&self,
sparse: impl IntoIterator<Item = (usize, T)>,
deg: usize,
) -> T
fn sparse_fold( &self, sparse: impl IntoIterator<Item = (usize, T)>, deg: usize, ) -> T
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 770)
731 pub fn solve_sparse_differential2(
732 p: &Self,
733 q: &Self,
734 x: &Self,
735 alpha: T,
736 beta: T,
737 deg: usize,
738 ) -> Self {
739 if deg == 0 {
740 return Self::zero();
741 }
742 let collect_sparse = |p: &Self| -> Vec<(usize, T)> {
743 p.iter()
744 .enumerate()
745 .filter(|&(_, x)| !x.is_zero())
746 .map(|(i, x)| (i, x.clone()))
747 .collect()
748 };
749 assert!(q.coeff(0).is_one());
750 assert!(x.coeff(0).is_one());
751 let p = collect_sparse(p);
752 let q = collect_sparse(q);
753 let x = collect_sparse(x);
754 let diff = |p: &[(usize, T)]| -> Vec<(usize, T)> {
755 p.iter()
756 .filter(|&&(i, _)| i > 0)
757 .map(|&(i, ref x)| (i - 1, x.clone() * T::from(i)))
758 .collect()
759 };
760 let dp = diff(&p);
761 let dq = diff(&q);
762
763 let mf = T::memorized_factorial(deg);
764 let mut f = Self::zeros(deg);
765 let mut qf = Self::zeros(deg);
766 let mut dq_f = Self::zeros(deg);
767 let mut d_qf = Self::zeros(deg);
768 f[0] = T::one();
769 for i in 0..deg - 1 {
770 qf[i] = f.sparse_fold(q.iter().cloned(), i);
771 dq_f[i] = f.sparse_fold(dq.iter().cloned(), i);
772 let dp_qf_i = qf.sparse_fold(dp.iter().cloned(), i);
773 let p_dq_f_i = dq_f.sparse_fold(p.iter().cloned(), i);
774 let x_d_qf_i = d_qf.sparse_fold(
775 x.iter()
776 .map(|&(i, ref x)| (i, x.clone() - T::from((i == 0) as usize))),
777 i,
778 );
779 d_qf[i] = alpha.clone() * dp_qf_i + beta.clone() * p_dq_f_i - x_d_qf_i;
780
781 let mut f_ip1 = d_qf[i].clone();
782 for &(j, ref q) in q.iter().take_while(|&&(j, _)| j <= i) {
783 if j > 0 {
784 f_ip1 -= q.clone() * &f[i - (j - 1)] * T::from(i - (j - 1));
785 }
786 }
787 f[i + 1] = f_ip1 * T::memorized_inv(&mf, i + 1);
788 }
789 f
790 }Sourcepub fn solve_sparse_differential2(
p: &Self,
q: &Self,
x: &Self,
alpha: T,
beta: T,
deg: usize,
) -> Self
pub fn solve_sparse_differential2( p: &Self, q: &Self, x: &Self, alpha: T, beta: T, deg: usize, ) -> Self
solve: $X(QF)’=\alpha P’(QF)+\beta P(Q’F)$ in $O(deg * max(nz(P), nz(Q), nz(X)))$
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (lines 839-846)
793 pub fn mul_of_pow_sparse(&self, q: &Self, exp_p: isize, exp_q: isize, deg: usize) -> Self {
794 if deg == 0 {
795 return Self::zero();
796 }
797 if exp_p == 0 && exp_q == 0 {
798 return Self::from_vec(
799 once(T::one())
800 .chain(repeat_with(T::zero))
801 .take(deg)
802 .collect(),
803 );
804 }
805 if exp_p != 0 && self.iter().all(|x| x.is_zero()) {
806 assert!(exp_p > 0);
807 return Self::zeros(deg);
808 }
809 if exp_q != 0 && q.iter().all(|x| x.is_zero()) {
810 assert!(exp_q > 0);
811 return Self::zeros(deg);
812 }
813
814 let normalize = |f: &Self, exp: isize| {
815 if exp == 0 {
816 return (0usize, T::one(), Self::from_vec(vec![T::one()]));
817 }
818 let k = f.iter().position(|value| !value.is_zero()).unwrap();
819 assert!(
820 exp >= 0 || k == 0,
821 "Negative exponent with zero constant term"
822 );
823 let c = f[k].clone();
824 let f = (f.clone() >> k) / &c;
825 (k, c, f)
826 };
827 let (sp, cp, mut p) = normalize(self, exp_p);
828 let (sq, cq, mut q) = normalize(q, exp_q);
829
830 let shift = exp_p
831 .saturating_mul(sp as _)
832 .saturating_add(exp_q.saturating_mul(sq as _)) as usize;
833 if shift >= deg {
834 return Self::zeros(deg);
835 }
836 p.truncate(deg - shift);
837 q.truncate(deg - shift);
838
839 let mut f = Self::solve_sparse_differential2(
840 &p,
841 &q,
842 &p,
843 T::from(exp_p),
844 T::from(exp_q),
845 deg - shift,
846 );
847 f *= cp.signed_pow(exp_p) * cq.signed_pow(exp_q);
848 if shift > 0 {
849 f <<= shift;
850 }
851 f.prefix(deg)
852 }
853
854 /// exp(P/Q)
855 pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self {
856 if deg == 0 {
857 return Self::zero();
858 }
859 let shift_q = q
860 .iter()
861 .position(|value| !value.is_zero())
862 .expect("Zero denominator");
863 let shift_p = self.iter().position(|value| !value.is_zero()).unwrap_or(!0);
864 assert!(shift_p > shift_q);
865
866 let mut p = self >> shift_q;
867 let mut q = q >> shift_q;
868 assert!(!q.coeff(0).is_zero());
869
870 let c = q[0].clone();
871 p /= c.clone();
872 q /= c;
873
874 Self::solve_sparse_differential2(&p, &q, &q, T::one(), -T::one(), deg)
875 }Sourcepub fn mul_of_pow_sparse(
&self,
q: &Self,
exp_p: isize,
exp_q: isize,
deg: usize,
) -> Self
pub fn mul_of_pow_sparse( &self, q: &Self, exp_p: isize, exp_q: isize, deg: usize, ) -> Self
P^exp_p * Q^exp_q
Sourcepub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self
pub fn exp_of_div_sparse(&self, q: &Self, deg: usize) -> Self
exp(P/Q)
Source§impl<T, C> FormalPowerSeries<T, C>
impl<T, C> FormalPowerSeries<T, C>
Sourcepub fn sqrt(&self, deg: usize) -> Option<Self>
pub fn sqrt(&self, deg: usize) -> Option<Self>
Examples found in repository?
More examples
crates/library_checker/src/polynomial/sqrt_of_formal_power_series_sparse.rs (line 14)
5pub fn sqrt_of_formal_power_series_sparse(reader: impl Read, writer: impl Write) {
6 prepare_io!(reader, writer);
7 sc!(n, k);
8 let mut a = vec![M::zero(); n];
9 for _ in 0..k {
10 sc!(i, a_i: M);
11 a[i] = a_i;
12 }
13 let f = Fps998244353::from_vec(a);
14 if let Some(g) = f.sqrt(n) {
15 pp!(@it g.data);
16 } else {
17 pp!("-1");
18 }
19}crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 889)
883 pub fn sqrt(&self, deg: usize) -> Option<Self> {
884 if self[0].is_zero() {
885 if let Some(k) = self.iter().position(|x| !x.is_zero()) {
886 if k % 2 != 0 {
887 return None;
888 } else if deg > k / 2 {
889 return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
890 }
891 }
892 } else {
893 let s = self[0].sqrt_coefficient()?;
894 if deg <= 1 {
895 return Some(Self::from(s).prefix(deg));
896 }
897 if let Some(step) = self.sparse_stride(deg, 4) {
898 let t = self[0].clone();
899 let mut f = self.prefix_ref(deg) / t;
900 f = f.pow_sparse1(T::one() / T::from(2usize), deg, step);
901 f *= s;
902 return Some(f);
903 }
904
905 let mut f = Self::from(s);
906 let inv2 = T::one() / (T::one() + T::one());
907 let inv2s = inv2.clone() / &f[0];
908 let extend = |f: &mut Self, end| {
909 for i in f.length()..end {
910 let mut value = self.coeff(i);
911 for j in 1..i {
912 value -= f[j].clone() * &f[i - j];
913 }
914 f.data.push(value * &inv2s);
915 }
916 };
917 extend(&mut f, deg.min(32));
918 f.truncate(deg);
919 if f.length() == deg {
920 return Some(f);
921 }
922 let mut inverse = f.inv(f.length());
923 let mut i = f.length();
924 while i < deg {
925 if deg - i <= 4 {
926 extend(&mut f, deg);
927 break;
928 }
929 let len = (i * 2).min(deg);
930 let factor = C::transform(inverse.data.clone(), i * 2);
931 let error = if !C::CYCLIC || i < 128 {
932 (self.prefix_ref(len) - &f * &f) >> i
933 } else {
934 let square = C::square(f.data.clone(), i);
935 // The cyclic square folds its high half into the already known low half.
936 Self::from_vec(
937 square
938 .into_iter()
939 .take(len - i)
940 .enumerate()
941 .map(|(j, value)| self.coeff(i + j) + self.coeff(j) - value)
942 .collect(),
943 )
944 };
945 let mut error_fft = C::transform(error.data, i * 2);
946 C::multiply(&mut error_fft, &factor);
947 let delta = C::inverse_transform(error_fft, i * 2);
948 f.data
949 .extend(delta.into_iter().take(len - i).map(|x| x * &inv2));
950 if i * 2 + 4 < deg {
951 let mut error_fft = C::transform(f.data.clone(), i * 2);
952 C::multiply(&mut error_fft, &factor);
953 let error = C::inverse_transform(error_fft, i * 2);
954 let mut error_fft = C::transform(error.into_iter().skip(i).collect(), i * 2);
955 C::multiply(&mut error_fft, &factor);
956 let error = C::inverse_transform(error_fft, i * 2);
957 inverse.data.extend(error.into_iter().take(i).map(Neg::neg));
958 }
959 i *= 2;
960 }
961 f.truncate(deg);
962 return Some(f);
963 }
964 Some(Self::zeros(deg))
965 }Source§impl<T, C> FormalPowerSeries<T, C>
impl<T, C> FormalPowerSeries<T, C>
Sourcepub fn count_subset_sum<F>(&self, deg: usize, inverse: F) -> Self
pub fn count_subset_sum<F>(&self, deg: usize, inverse: F) -> Self
pub fn count_multiset_sum<F>(&self, deg: usize, inverse: F) -> Self
Sourcepub fn bostan_mori(self, rhs: Self, n: usize) -> T
pub fn bostan_mori(self, rhs: Self, n: usize) -> T
[x^n] P(x) / Q(x)
Sourcepub fn bostan_mori_msb(self, n: usize) -> Self
pub fn bostan_mori_msb(self, n: usize) -> Self
return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 1064)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }Sourcepub fn pow_mod(self, n: usize) -> Self
pub fn pow_mod(self, n: usize) -> Self
x^n mod self
Examples found in repository?
crates/competitive/src/math/black_box_mint_matrix.rs (line 42)
33 fn apply_pow<C>(&self, mut b: Vec<MInt<M>>, k: usize) -> Vec<MInt<M>>
34 where
35 C: ConvolveSteps<T = Vec<MInt<M>>>,
36 {
37 assert_eq!(self.shape().0, self.shape().1);
38 assert_eq!(self.shape().1, b.len());
39 let n = self.shape().0;
40 let p = self.minimal_polynomial();
41 let polynomial: FormalPowerSeries<MInt<M>, C> = FormalPowerSeries::from_vec(p);
42 let f = polynomial.pow_mod(k);
43 let mut res = vec![MInt::zero(); n];
44 for f in f {
45 for j in 0..n {
46 res[j] += f * b[j];
47 }
48 b = self.apply(&b);
49 }
50 res
51 }Sourcefn middle_product(self, other: &C::F, deg: usize) -> Self
fn middle_product(self, other: &C::F, deg: usize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 1071)
1057 pub fn bostan_mori_msb(self, n: usize) -> Self {
1058 let d = self.length() - 1;
1059 if n == 0 {
1060 return (Self::one() << (d - 1)) / self[0].clone();
1061 }
1062 let q = self;
1063 let mq = q.clone().parity_inversion();
1064 let w = (q * &mq).even().bostan_mori_msb(n / 2);
1065 let mut s = Self::zeros(w.length() * 2 - (n % 2));
1066 for (i, x) in w.iter().enumerate() {
1067 s[i * 2 + (1 - n % 2)] = x.clone();
1068 }
1069 let len = 2 * d + 1;
1070 let ts = C::transform(s.prefix(len).data, len);
1071 mq.reversed().middle_product(&ts, len).prefix(d + 1)
1072 }
1073 /// x^n mod self
1074 pub fn pow_mod(self, n: usize) -> Self {
1075 let d = self.length() - 1;
1076 let q = self.reversed();
1077 let u = q.clone().bostan_mori_msb(n);
1078 let mut f = (u * q).prefix(d).reversed();
1079 f.trim_tail_zeros();
1080 f
1081 }
1082 fn middle_product(self, other: &C::F, deg: usize) -> Self {
1083 let n = self.length();
1084 let mut s = C::transform(self.reversed().data, deg);
1085 C::multiply(&mut s, other);
1086 Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
1087 }
1088 pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
1089 where
1090 C: NttReuse<T = Vec<T>>,
1091 C::F: Clone,
1092 {
1093 let n = points.len();
1094 if n <= 32 || self.length() <= 32 {
1095 return points.iter().map(|p| self.eval(p.clone())).collect();
1096 }
1097 let size = n.next_power_of_two();
1098 let block = 16;
1099 let leaves = size / block;
1100 let mut subproduct_tree = Vec::with_capacity(leaves * 2);
1101 subproduct_tree.resize_with(leaves * 2, || None);
1102 let mut leaf_products = Vec::with_capacity(leaves);
1103 for i in 0..leaves {
1104 let mut product = vec![T::one()];
1105 for j in 0..block {
1106 let x = points.get(i * block + j).cloned().unwrap_or_else(T::zero);
1107 product.push(T::one());
1108 for k in (1..=j).rev() {
1109 product[k] = product[k - 1].clone() - x.clone() * &product[k];
1110 }
1111 product[0] *= -x;
1112 }
1113 subproduct_tree[leaves + i] = Some(C::transform_ntt(product.clone(), block * 2));
1114 leaf_products.push(product);
1115 }
1116 for i in (1..leaves).rev() {
1117 let mut product = subproduct_tree[i * 2].as_ref().unwrap().clone();
1118 C::multiply_prefix(&mut product, subproduct_tree[i * 2 + 1].as_ref().unwrap());
1119 if i > 1 {
1120 product = C::ntt_doubling(product, true);
1121 }
1122 subproduct_tree[i] = Some(product);
1123 }
1124 let mut product = C::inverse_transform_ntt(subproduct_tree[1].take().unwrap(), size);
1125 product[0] -= T::one();
1126 product.push(T::one());
1127 let mut uptree_t = Vec::with_capacity(leaves * 2);
1128 uptree_t.resize_with(1, Zero::zero);
1129 let m = self.length();
1130 let v = Self::from_vec(product).reversed().resized(m);
1131 let s = C::transform(self.data, m * 2);
1132 uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(size));
1133 for i in 1..leaves {
1134 let degree = uptree_t[i].length();
1135 let spectrum = C::transform_ntt(std::mem::take(&mut uptree_t[i].data), degree);
1136 let left = subproduct_tree[i * 2].take().unwrap();
1137 let right = subproduct_tree[i * 2 + 1].take().unwrap();
1138 let mut child = spectrum.clone();
1139 C::multiply_prefix(&mut child, &right);
1140 let mut child = C::inverse_transform_ntt(child, degree);
1141 child.drain(..degree / 2);
1142 uptree_t.push(Self::from_vec(child));
1143 let mut child = spectrum;
1144 C::multiply_prefix(&mut child, &left);
1145 let mut child = C::inverse_transform_ntt(child, degree);
1146 child.drain(..degree / 2);
1147 uptree_t.push(Self::from_vec(child));
1148 }
1149 let mut result = Vec::with_capacity(n);
1150 for ((values, product), points) in uptree_t[leaves..]
1151 .iter()
1152 .zip(leaf_products)
1153 .zip(points.chunks(block))
1154 {
1155 let mut remainder = Self::zeros(block);
1156 for (j, value) in values.iter().enumerate() {
1157 for (r, p) in remainder.data[..=j].iter_mut().zip(&product[block - j..]) {
1158 *r += value.clone() * p;
1159 }
1160 }
1161 result.extend(points.iter().map(|p| remainder.eval(p.clone())));
1162 }
1163 result
1164 }Sourcepub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T>
Sourcepub fn product_all<I>(iter: I, deg: usize) -> Selfwhere
I: IntoIterator<Item = Self>,
pub fn product_all<I>(iter: I, deg: usize) -> Selfwhere
I: IntoIterator<Item = Self>,
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (lines 1248-1254)
1243 pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
1244 where
1245 I: IntoIterator<Item = T>,
1246 {
1247 let mut n = T::zero();
1248 let prod = Self::product_all(
1249 iter.into_iter().map(|a| {
1250 n += T::one();
1251 Self::from_vec(vec![T::one(), -a])
1252 }),
1253 deg,
1254 );
1255 (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
1256 }Sourcepub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)where
I: IntoIterator<Item = (Self, Self)>,
pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)where
I: IntoIterator<Item = (Self, Self)>,
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (lines 1231-1235)
1226 pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
1227 where
1228 I: IntoIterator<Item = (T, T)>,
1229 F: FnMut(usize) -> T,
1230 {
1231 let (p, q) = Self::sum_all_rational(
1232 iter.into_iter()
1233 .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
1234 deg,
1235 );
1236 let mut f = (p * q.inv(deg)).prefix(deg);
1237 for i in 0..f.length() {
1238 f[i] *= inv_fact(i);
1239 }
1240 f
1241 }Sourcepub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
Examples found in repository?
More examples
pub fn kth_term(a: Vec<T>, k: usize) -> T
Sourcepub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, inv_fact: F) -> Self
pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, inv_fact: F) -> Self
sum_i a_i exp(b_i x)
Sourcepub fn sum_of_powers<I>(iter: I, deg: usize) -> Selfwhere
I: IntoIterator<Item = T>,
pub fn sum_of_powers<I>(iter: I, deg: usize) -> Selfwhere
I: IntoIterator<Item = T>,
sum_i (a_i x)^j
Sourcepub fn power_projection(&self, w: &[T], m: usize) -> Self
pub fn power_projection(&self, w: &[T], m: usize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 1343)
1322 pub fn compositional_inverse(&self, deg: usize) -> Self
1323 where
1324 C: NttReuse<T = Vec<T>>,
1325 C::F: Clone,
1326 {
1327 if deg == 0 {
1328 return Self::zero();
1329 }
1330 if deg == 1 {
1331 return Self::from_vec(vec![T::zero()]);
1332 }
1333 debug_assert!(self[0].is_zero());
1334 debug_assert!(!self[1].is_zero());
1335
1336 let mut f = self.prefix_ref(deg);
1337 f.resize(deg);
1338 let c = f[1].clone();
1339 f /= c.clone();
1340
1341 let mut w = vec![T::zero(); deg];
1342 w[deg - 1] = T::one();
1343 let s = f.power_projection(&w, deg);
1344
1345 let n = deg - 1;
1346 let n_t = T::from(n);
1347 let mut h = vec![T::zero(); n];
1348 for i in 1..=n {
1349 h[n - i] = s[i].clone() * &n_t / T::from(i);
1350 }
1351
1352 let h_fps = Self::from_vec(h);
1353 let inv_n = T::one() / n_t;
1354 let mut t = h_fps.log(n);
1355 t *= -inv_n;
1356 let g_over_x = t.exp(n);
1357 let mut g = (g_over_x << 1).prefix(deg);
1358
1359 let inv_c = T::one() / c;
1360 let mut pow = T::one();
1361 for coef in g.iter_mut() {
1362 *coef *= pow.clone();
1363 pow *= inv_c.clone();
1364 }
1365 g
1366 }Sourcepub fn compositional_inverse(&self, deg: usize) -> Self
pub fn compositional_inverse(&self, deg: usize) -> Self
Examples found in repository?
More examples
Source§impl<T, C> FormalPowerSeries<T, C>
impl<T, C> FormalPowerSeries<T, C>
Sourcepub fn div_rem(self, rhs: Self) -> (Self, Self)
pub fn div_rem(self, rhs: Self) -> (Self, Self)
Examples found in repository?
More examples
crates/competitive/src/math/formal_power_series/berlekamp_massey.rs (line 252)
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 }
265}
266
267fn degree<T, C>(fps: &FormalPowerSeries<T, C>) -> isize {
268 fps.length() as isize - 1
269}
270
271fn add<T, C>(
272 left: FormalPowerSeries<T, C>,
273 right: FormalPowerSeries<T, C>,
274) -> FormalPowerSeries<T, C>
275where
276 T: FormalPowerSeriesCoefficient,
277{
278 (left + right).trimed()
279}
280
281fn tail<T, C>(fps: &FormalPowerSeries<T, C>, start: isize) -> FormalPowerSeries<T, C>
282where
283 T: FormalPowerSeriesCoefficient,
284{
285 let start = start.max(0) as usize;
286 if start >= fps.length() {
287 FormalPowerSeries::zero()
288 } else {
289 FormalPowerSeries::from_vec(fps.data[start..].to_vec())
290 }
291}
292
293fn coefficient<T, C>(fps: &FormalPowerSeries<T, C>, index: isize) -> T
294where
295 T: FormalPowerSeriesCoefficient,
296{
297 if index < 0 {
298 T::zero()
299 } else {
300 fps.coeff(index as usize)
301 }
302}
303
304fn brute_force<T, C>(
305 mut p: FormalPowerSeries<T, C>,
306 mut q: FormalPowerSeries<T, C>,
307 k: usize,
308) -> FpsMatrix<T, C>
309where
310 T: FormalPowerSeriesCoefficient,
311 C: NttReuse<T = Vec<T>>,
312 C::F: Clone,
313{
314 let threshold = degree(&p) - k as isize;
315 let mut matrix = FpsMatrix::identity();
316 while q.length() as isize > threshold {
317 let q_degree = q.length() - 1;
318 let mut negative_quotient = vec![T::zero(); p.length() - q.length() + 1];
319 let inverse = -T::one() / &q[q_degree];
320 for i in (0..negative_quotient.len()).rev() {
321 negative_quotient[i] = p[i + q_degree].clone() * &inverse;
322 p[i + q_degree] = T::zero();
323 for j in 0..q_degree {
324 let value = negative_quotient[i].clone() * &q[j];
325 p[i + j] += &value;
326 }
327 }
328 matrix.left_multiply_step(&negative_quotient);
329 p.truncate(q_degree);
330 p.trim_tail_zeros();
331 swap(&mut p, &mut q);
332 }
333 matrix
334}
335
336fn reduced_transform<T, C>(fps: &FormalPowerSeries<T, C>, length: usize) -> C::F
337where
338 T: FormalPowerSeriesCoefficient,
339 C: NttReuse<T = Vec<T>>,
340{
341 let mut coefficients = vec![T::zero(); length];
342 for (i, value) in fps.iter().enumerate() {
343 coefficients[i & (length - 1)] += value;
344 }
345 C::transform_ntt(coefficients, length)
346}
347
348fn transform_window<T, C>(fps: &FormalPowerSeries<T, C>, end: isize, length: usize) -> C::F
349where
350 T: FormalPowerSeriesCoefficient,
351 C: NttReuse<T = Vec<T>>,
352{
353 let start = end - length as isize;
354 let coefficients = (start..end).map(|index| coefficient(fps, index)).collect();
355 C::transform_ntt(coefficients, length)
356}
357
358fn half_gcd<T, C>(
359 p: &FormalPowerSeries<T, C>,
360 q: &FormalPowerSeries<T, C>,
361 k: usize,
362 length: usize,
363) -> (FpsMatrix<T, C>, FrequencyMatrix<C>)
364where
365 T: FormalPowerSeriesCoefficient,
366 C: NttReuse<T = Vec<T>>,
367 C::F: Clone,
368{
369 let d = degree(p);
370 if degree(q) < d - k as isize {
371 let matrix = FpsMatrix::identity();
372 let frequency = matrix.transform(length);
373 return (matrix, frequency);
374 }
375 if k == 1 {
376 let matrix = FpsMatrix {
377 a00: FormalPowerSeries::zero(),
378 a01: FormalPowerSeries::one(),
379 a10: FormalPowerSeries::one(),
380 a11: -(tail(p, d - 2) / tail(q, d - 2)),
381 };
382 let frequency = matrix.transform(length);
383 return (matrix, frequency);
384 }
385 if p.length().min(q.length()) <= 32 {
386 let matrix = brute_force(p.clone(), q.clone(), k);
387 let frequency = matrix.transform(length);
388 return (matrix, frequency);
389 }
390
391 let half = length / 2;
392 if k <= half {
393 let (matrix, frequency) = half_gcd(p, q, k, half);
394 let frequency = matrix.extend_transform(frequency, length);
395 return (matrix, frequency);
396 }
397
398 let (matrix, mut matrix_frequency) = half_gcd(
399 &tail(p, d - 2 * half as isize),
400 &tail(q, d - 2 * half as isize),
401 half,
402 length,
403 );
404 let degeneracy = half as isize - degree(&matrix.a11);
405
406 let (p0, q0) = matrix_frequency.apply(
407 &transform_window(p, d - half as isize + degeneracy, length),
408 &transform_window(q, d - half as isize + degeneracy, length),
409 length,
410 );
411 let (p1, q1) = matrix_frequency.apply(
412 &transform_window(p, d - 2 * half as isize, length),
413 &transform_window(q, d - 2 * half as isize, length),
414 length,
415 );
416 let part_length = (half as isize + degeneracy) as usize;
417 let mut p_reduced = p1[length - part_length..].to_vec();
418 p_reduced.extend_from_slice(&p0[length - part_length..]);
419 let mut q_reduced = q1[length - part_length..].to_vec();
420 q_reduced.extend_from_slice(&q0[length - part_length..]);
421 let mut q_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(q_reduced).trimed();
422
423 let position = d - half as isize + degeneracy;
424 let mut leading = T::zero();
425 for i in 0..=position {
426 leading += coefficient(p, i) * coefficient(&matrix.a00, position - i)
427 + coefficient(q, i) * coefficient(&matrix.a01, position - i);
428 }
429 p_reduced.push(leading);
430 let mut p_reduced: FormalPowerSeries<T, C> = FormalPowerSeries::from_vec(p_reduced);
431 if degree(&q_reduced) < 3 * half as isize + degeneracy - k as isize {
432 return (matrix, matrix_frequency);
433 }
434
435 let mut remaining = k as isize - degree(&matrix.a11);
436 let mut top_product = matrix.a11.data.last().unwrap().clone();
437 let mut product_degree = degree(&matrix.a11);
438 if degeneracy > 0 {
439 let skip = (2 * half as isize + 2 * degeneracy - (d - half as isize + degeneracy)).max(0);
440 let (division, remainder) = tail(&p_reduced, skip).div_rem(tail(&q_reduced, skip));
441 remaining -= degree(&division);
442 top_product *= -division.data.last().unwrap().clone();
443 product_degree += degree(&division);
444 matrix_frequency = matrix_frequency.left_multiply_step(&division, length);
445 swap(&mut p_reduced, &mut q_reduced);
446 q_reduced = FormalPowerSeries::zeros(skip as usize);
447 q_reduced.data.extend(remainder.data);
448 }
449
450 let start = 3 * half as isize + degeneracy - k as isize - remaining;
451 let (right_matrix, right_frequency) = half_gcd(
452 &tail(&p_reduced, start),
453 &tail(&q_reduced, start),
454 remaining as usize,
455 length,
456 );
457 let product_frequency = right_frequency.multiply(&matrix_frequency);
458 let mut product = product_frequency.clone().inverse_transform(length);
459 product.a00.truncate(k);
460 product.a00.trim_tail_zeros();
461 product.a01.truncate(k);
462 product.a01.trim_tail_zeros();
463 product.a10.truncate(k);
464 product.a10.trim_tail_zeros();
465 product_degree += degree(&right_matrix.a11);
466 if product_degree == length as isize {
467 product.a11.resize(k + 1);
468 let highest = top_product * right_matrix.a11.data.last().unwrap();
469 product.a11[k] = highest.clone();
470 product.a11[0] -= highest;
471 }
472 product.a11.trim_tail_zeros();
473 let product_frequency = if C::MULTIPLE {
474 product.transform(length)
475 } else {
476 product_frequency
477 };
478 (product, product_frequency)
479}Trait Implementations§
Source§impl<T, C> Add for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Add for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Add<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Add<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
+ operator.Source§impl<T, C> Add<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Add<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
+ operator.Source§impl<T, C> Add<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Add<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Add<&T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Add<&T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Add<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Add<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
+ operator.Source§impl<T, C> Add<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Add<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Add<T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Add<T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> AddAssign for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> AddAssign for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn add_assign(&mut self, rhs: Self)
fn add_assign(&mut self, rhs: Self)
Performs the
+= operation. Read moreSource§impl<T, C> AddAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> AddAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn add_assign(&mut self, rhs: &Self)
fn add_assign(&mut self, rhs: &Self)
Performs the
+= operation. Read moreSource§impl<T, C> AddAssign<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> AddAssign<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn add_assign(&mut self, rhs: &T)
fn add_assign(&mut self, rhs: &T)
Performs the
+= operation. Read moreSource§impl<T, C> AddAssign<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> AddAssign<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn add_assign(&mut self, rhs: T)
fn add_assign(&mut self, rhs: T)
Performs the
+= operation. Read moreSource§impl<T, C> Clone for FormalPowerSeries<T, C>where
T: Clone,
impl<T, C> Clone for FormalPowerSeries<T, C>where
T: Clone,
Source§impl<T, C> Debug for FormalPowerSeries<T, C>where
T: Debug,
impl<T, C> Debug for FormalPowerSeries<T, C>where
T: Debug,
Source§impl<T, C> Div for FormalPowerSeries<T, C>
impl<T, C> Div for FormalPowerSeries<T, C>
Source§impl<T, C> Div<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
impl<T, C> Div<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
/ operator.Source§impl<T, C> Div<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
impl<T, C> Div<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
/ operator.Source§impl<T, C> Div<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Div<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Div<&T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Div<&T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Div<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
impl<T, C> Div<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
/ operator.Source§impl<T, C> Div<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Div<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Div<T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Div<T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> DivAssign for FormalPowerSeries<T, C>
impl<T, C> DivAssign for FormalPowerSeries<T, C>
Source§fn div_assign(&mut self, rhs: Self)
fn div_assign(&mut self, rhs: Self)
Performs the
/= operation. Read moreSource§impl<T, C> DivAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
impl<T, C> DivAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
Source§fn div_assign(&mut self, rhs: &Self)
fn div_assign(&mut self, rhs: &Self)
Performs the
/= operation. Read moreSource§impl<T, C> DivAssign<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> DivAssign<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn div_assign(&mut self, rhs: &T)
fn div_assign(&mut self, rhs: &T)
Performs the
/= operation. Read moreSource§impl<T, C> DivAssign<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> DivAssign<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn div_assign(&mut self, rhs: T)
fn div_assign(&mut self, rhs: T)
Performs the
/= operation. Read moreimpl<T, C> Eq for FormalPowerSeries<T, C>where
T: PartialEq,
Source§impl<T, C> From<T> for FormalPowerSeries<T, C>
impl<T, C> From<T> for FormalPowerSeries<T, C>
Source§impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C>
impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C>
Source§impl<T, C> FromIterator<T> for FormalPowerSeries<T, C>
impl<T, C> FromIterator<T> for FormalPowerSeries<T, C>
Source§fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self
Creates a value from an iterator. Read more
Source§impl<T, C> Index<usize> for FormalPowerSeries<T, C>
impl<T, C> Index<usize> for FormalPowerSeries<T, C>
Source§impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C>
impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C>
Source§impl<T, C> IntoIterator for FormalPowerSeries<T, C>
impl<T, C> IntoIterator for FormalPowerSeries<T, C>
Source§impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C>
impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C>
Source§impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C>
impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C>
Source§impl<T, C> Mul for FormalPowerSeries<T, C>where
C: ConvolveSteps<T = Vec<T>>,
impl<T, C> Mul for FormalPowerSeries<T, C>where
C: ConvolveSteps<T = Vec<T>>,
Source§impl<T, C> Mul<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
impl<T, C> Mul<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
* operator.Source§impl<T, C> Mul<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
impl<T, C> Mul<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
* operator.Source§impl<T, C> Mul<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Mul<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Mul<&T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Mul<&T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Mul<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
impl<T, C> Mul<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
* operator.Source§impl<T, C> Mul<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Mul<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Mul<T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Mul<T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> MulAssign for FormalPowerSeries<T, C>
impl<T, C> MulAssign for FormalPowerSeries<T, C>
Source§fn mul_assign(&mut self, rhs: Self)
fn mul_assign(&mut self, rhs: Self)
Performs the
*= operation. Read moreSource§impl<T, C> MulAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
impl<T, C> MulAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
Source§fn mul_assign(&mut self, rhs: &Self)
fn mul_assign(&mut self, rhs: &Self)
Performs the
*= operation. Read moreSource§impl<T, C> MulAssign<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> MulAssign<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn mul_assign(&mut self, rhs: &T)
fn mul_assign(&mut self, rhs: &T)
Performs the
*= operation. Read moreSource§impl<T, C> MulAssign<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> MulAssign<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn mul_assign(&mut self, rhs: T)
fn mul_assign(&mut self, rhs: T)
Performs the
*= operation. Read moreSource§impl<T, C> Neg for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Neg for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Neg for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Neg for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> One for FormalPowerSeries<T, C>
impl<T, C> One for FormalPowerSeries<T, C>
Source§impl<T, C> PartialEq for FormalPowerSeries<T, C>where
T: PartialEq,
impl<T, C> PartialEq for FormalPowerSeries<T, C>where
T: PartialEq,
Source§impl<T, C> Rem for FormalPowerSeries<T, C>
impl<T, C> Rem for FormalPowerSeries<T, C>
Source§impl<T, C> Rem<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
impl<T, C> Rem<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
% operator.Source§impl<T, C> Rem<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
impl<T, C> Rem<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
% operator.Source§impl<T, C> Rem<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
impl<T, C> Rem<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
% operator.Source§impl<T, C> RemAssign for FormalPowerSeries<T, C>
impl<T, C> RemAssign for FormalPowerSeries<T, C>
Source§fn rem_assign(&mut self, rhs: Self)
fn rem_assign(&mut self, rhs: Self)
Performs the
%= operation. Read moreSource§impl<T, C> RemAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
impl<T, C> RemAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
Source§fn rem_assign(&mut self, rhs: &Self)
fn rem_assign(&mut self, rhs: &Self)
Performs the
%= operation. Read moreSource§impl<T, C> Shl<usize> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Shl<usize> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Shl<usize> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Shl<usize> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> ShlAssign<usize> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> ShlAssign<usize> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn shl_assign(&mut self, rhs: usize)
fn shl_assign(&mut self, rhs: usize)
Performs the
<<= operation. Read moreSource§impl<T, C> Shr<usize> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Shr<usize> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Shr<usize> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Shr<usize> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> ShrAssign<usize> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> ShrAssign<usize> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn shr_assign(&mut self, rhs: usize)
fn shr_assign(&mut self, rhs: usize)
Performs the
>>= operation. Read moreSource§impl<T, C> Sub for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Sub for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Sub<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Sub<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
- operator.Source§impl<T, C> Sub<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Sub<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
- operator.Source§impl<T, C> Sub<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Sub<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Sub<&T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Sub<&T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Sub<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Sub<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§type Output = FormalPowerSeries<T, C>
type Output = FormalPowerSeries<T, C>
The resulting type after applying the
- operator.Source§impl<T, C> Sub<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Sub<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> Sub<T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> Sub<T> for &FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§impl<T, C> SubAssign for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> SubAssign for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn sub_assign(&mut self, rhs: Self)
fn sub_assign(&mut self, rhs: Self)
Performs the
-= operation. Read moreSource§impl<T, C> SubAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> SubAssign<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn sub_assign(&mut self, rhs: &Self)
fn sub_assign(&mut self, rhs: &Self)
Performs the
-= operation. Read moreSource§impl<T, C> SubAssign<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> SubAssign<&T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn sub_assign(&mut self, rhs: &T)
fn sub_assign(&mut self, rhs: &T)
Performs the
-= operation. Read moreSource§impl<T, C> SubAssign<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
impl<T, C> SubAssign<T> for FormalPowerSeries<T, C>where
T: FormalPowerSeriesCoefficient,
Source§fn sub_assign(&mut self, rhs: T)
fn sub_assign(&mut self, rhs: T)
Performs the
-= operation. Read moreAuto Trait Implementations§
impl<T, C> Freeze for FormalPowerSeries<T, C>
impl<T, C> RefUnwindSafe for FormalPowerSeries<T, C>
impl<T, C> Send for FormalPowerSeries<T, C>
impl<T, C> Sync for FormalPowerSeries<T, C>
impl<T, C> Unpin for FormalPowerSeries<T, C>
impl<T, C> UnsafeUnpin for FormalPowerSeries<T, C>
impl<T, C> UnwindSafe for FormalPowerSeries<T, C>
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more