pub trait FormalPowerSeriesCoefficient:
Sized
+ Clone
+ Zero
+ PartialEq
+ One
+ From<usize>
+ From<isize>
+ Add<Output = Self>
+ Sub<Output = Self>
+ Mul<Output = Self>
+ Div<Output = Self>
+ for<'r> Add<&'r Self, Output = Self>
+ for<'r> Sub<&'r Self, Output = Self>
+ for<'r> Mul<&'r Self, Output = Self>
+ for<'r> Div<&'r Self, Output = Self>
+ AddAssign<Self>
+ SubAssign<Self>
+ MulAssign<Self>
+ DivAssign<Self>
+ for<'r> AddAssign<&'r Self>
+ for<'r> SubAssign<&'r Self>
+ for<'r> MulAssign<&'r Self>
+ for<'r> DivAssign<&'r Self>
+ Neg<Output = Self> {
type Base: MIntConvert<usize> + MIntConvert<isize>;
// Required methods
fn pow(self, exp: usize) -> Self;
fn memorized_fact(mf: &MemorizedFactorial<Self::Base>) -> &[Self];
fn memorized_inv_fact(mf: &MemorizedFactorial<Self::Base>) -> &[Self];
fn memorized_inv(mf: &MemorizedFactorial<Self::Base>, n: usize) -> Self;
// Provided methods
fn signed_pow(self, exp: isize) -> Self { ... }
fn memorized_factorial(n: usize) -> MemorizedFactorial<Self::Base> { ... }
}Required Associated Types§
type Base: MIntConvert<usize> + MIntConvert<isize>
Required Methods§
fn pow(self, exp: usize) -> Self
fn memorized_fact(mf: &MemorizedFactorial<Self::Base>) -> &[Self]
fn memorized_inv_fact(mf: &MemorizedFactorial<Self::Base>) -> &[Self]
fn memorized_inv(mf: &MemorizedFactorial<Self::Base>, n: usize) -> Self
Provided Methods§
Sourcefn signed_pow(self, exp: isize) -> Self
fn signed_pow(self, exp: isize) -> Self
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 847)
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 }Sourcefn memorized_factorial(n: usize) -> MemorizedFactorial<Self::Base>
fn memorized_factorial(n: usize) -> MemorizedFactorial<Self::Base>
Examples found in repository?
crates/competitive/src/math/formal_power_series/formal_power_series_impls.rs (line 357)
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 }Dyn Compatibility§
This trait is not dyn compatible.
In older versions of Rust, dyn compatibility was called "object safety".