From 4d4e1a8e0c928e863071dbcf89e06652b2c56203 Mon Sep 17 00:00:00 2001 From: Nagata Kana Date: Tue, 28 May 2024 21:08:31 +0900 Subject: [PATCH] Add naive_poly --- libs/{poly => naive_poly}/Cargo.toml | 6 +- libs/naive_poly/src/lib.rs | 435 +++++++++++++++++++++++++++ libs/poly/src/lib.rs | 303 ------------------- 3 files changed, 438 insertions(+), 306 deletions(-) rename libs/{poly => naive_poly}/Cargo.toml (73%) create mode 100644 libs/naive_poly/src/lib.rs delete mode 100644 libs/poly/src/lib.rs diff --git a/libs/poly/Cargo.toml b/libs/naive_poly/Cargo.toml similarity index 73% rename from libs/poly/Cargo.toml rename to libs/naive_poly/Cargo.toml index 31839910..40dd86aa 100644 --- a/libs/poly/Cargo.toml +++ b/libs/naive_poly/Cargo.toml @@ -1,11 +1,11 @@ [package] -name = "poly" +name = "naive_poly" version = "0.1.0" -edition = "2018" +edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] [dev-dependencies] -rand = { workspace = true } +rand.workspace = true diff --git a/libs/naive_poly/src/lib.rs b/libs/naive_poly/src/lib.rs new file mode 100644 index 00000000..549f35a2 --- /dev/null +++ b/libs/naive_poly/src/lib.rs @@ -0,0 +1,435 @@ +//! # Naive implementation of polynomial operations + +use std::iter::Product; +use std::iter::Sum; +use std::ops::AddAssign; +use std::ops::Div; +use std::ops::Mul; +use std::ops::MulAssign; +use std::ops::SubAssign; + +fn zero() -> T +where + T: Sum, +{ + std::iter::empty().sum() +} + +fn one() -> T +where + T: Product, +{ + std::iter::empty().product() +} + +/// Multiply two polynomials in $O(nm)$ time. +/// +/// If $a = 0$ or $b = 0$, returns $0$ (empty vector). +pub fn mul(a: &[T], b: &[T]) -> Vec +where + T: Copy + Sum + Mul + AddAssign, +{ + if a.is_empty() || b.is_empty() { + return vec![zero(); 0]; + } + let mut c = vec![zero(); a.len() + b.len() - 1]; + for (i, &ai) in a.iter().enumerate() { + for (j, &bj) in b.iter().enumerate() { + c[i + j] += ai * bj; + } + } + c +} + +/// Add two polynomials in $O(\max(n, m))$ time. +/// +/// Trailing zeros are removed. +pub fn add(a: &[T], b: &[T]) -> Vec +where + T: Copy + PartialEq + Sum + AddAssign, +{ + let mut c = vec![zero(); a.len().max(b.len())]; + for (i, &ai) in a.iter().enumerate() { + c[i] += ai; + } + for (i, &bi) in b.iter().enumerate() { + c[i] += bi; + } + while c.last() == Some(&zero()) { + c.pop().unwrap(); + } + c +} + +/// Subtract two polynomials in $O(\max(n, m))$ time. +/// +/// Trailing zeros are removed. +pub fn sub(a: &[T], b: &[T]) -> Vec +where + T: Copy + PartialEq + Sum + AddAssign + SubAssign + Mul, +{ + let mut c = vec![zero(); a.len().max(b.len())]; + for (i, &ai) in a.iter().enumerate() { + c[i] += ai; + } + for (i, &bi) in b.iter().enumerate() { + c[i] -= bi; + } + while c.last() == Some(&zero()) { + c.pop().unwrap(); + } + c +} + +/// Divide two polynomials in $O((n - m) m)$ time. +/// +/// * Returns the quotient and modifies $a$ to store the remainder. +/// * Trailing zeros in $a$ are removed from the remainder. +/// * $b_{m-1} \neq 0$ is required. +pub fn div(a: &mut Vec, b: &[T]) -> Vec +where + T: Copy + PartialEq + Sum + SubAssign + Mul + Div, +{ + assert!(!b.is_empty() && *b.last().unwrap() != zero::()); + if a.len() < b.len() { + return vec![zero(); 0]; + } + let mut c = vec![zero(); a.len() + 1 - b.len()]; + for i in (0..c.len()).rev() { + c[i] = a[i] / *b.last().unwrap(); + for j in 0..b.len() { + a[i + j] -= c[i] * b[j]; + } + } + while a.last() == Some(&zero()) { + a.pop().unwrap(); + } + c +} + +/// Evaluate a polynomial at a point in $O(n)$ time. +pub fn eval(p: &[T], x: T) -> T +where + T: Copy + Sum + Mul + AddAssign + MulAssign, +{ + let mut y = zero(); + for &pi in p.iter().rev() { + y *= x; + y += pi; + } + y +} + +/// Compute the $e$-th power of a polynomial in $O((en)^2 \log e)$ time. +pub fn pow(mut a: Vec, mut e: u64) -> Vec +where + T: Copy + PartialEq + Sum + Product + AddAssign + Mul, +{ + if e == 0 { + return vec![one(); 1]; + } + let mut b = vec![zero(); 1]; + b[0] = one(); + while e > 1 { + if e % 2 == 1 { + b = mul(&b, &a); + } + a = mul(&a, &a); + e /= 2; + } + mul(&b, &a) +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::rngs::StdRng; + use rand::Rng; + use rand::SeedableRng; + use std::ops::Add; + use std::ops::DivAssign; + use std::ops::MulAssign; + use std::ops::Sub; + + #[derive(Debug, Clone, Copy, PartialEq)] + struct Fp(u64); + impl Fp

{ + fn new(x: u64) -> Self { + Self(x % P) + } + } + impl AddAssign for Fp

{ + fn add_assign(&mut self, rhs: Self) { + self.0 += rhs.0; + if self.0 >= P { + self.0 -= P; + } + } + } + impl Add for Fp

{ + type Output = Self; + + fn add(mut self, rhs: Self) -> Self { + self += rhs; + self + } + } + impl SubAssign for Fp

{ + fn sub_assign(&mut self, rhs: Self) { + if self.0 < rhs.0 { + self.0 += P; + } + self.0 -= rhs.0; + } + } + impl Sub for Fp

{ + type Output = Self; + + fn sub(mut self, rhs: Self) -> Self { + self -= rhs; + self + } + } + impl Mul for Fp

{ + type Output = Self; + + fn mul(self, rhs: Self) -> Self { + Self::new(self.0 * rhs.0 % P) + } + } + impl MulAssign for Fp

{ + fn mul_assign(&mut self, rhs: Self) { + *self = *self * rhs; + } + } + impl Div for Fp

{ + type Output = Self; + + fn div(mut self, rhs: Self) -> Self { + self /= rhs; + self + } + } + impl DivAssign for Fp

{ + fn div_assign(&mut self, mut rhs: Self) { + let mut e = P - 2; + while e > 0 { + if e % 2 == 1 { + *self *= rhs; + } + rhs *= rhs; + e /= 2; + } + } + } + impl Sum> for Fp

{ + fn sum>>(iter: I) -> Self { + iter.fold(Fp::

(0), |a, b| a + b) + } + } + + fn gen_poly(n: usize, rng: &mut StdRng) -> Vec> { + (0..n) + .map(|i| { + if i + 1 == n { + Fp::

(rng.gen_range(1..P)) + } else { + Fp::

(rng.gen_range(0..P)) + } + }) + .collect() + } + + #[test] + fn test_add() { + const P: u64 = 13; + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..1000 { + let n = rng.gen_range(0..10); + let a = gen_poly::

(n, &mut rng); + let b = gen_poly::

(n, &mut rng); + let c = add(&a, &b); + assert!(c.len() <= a.len().max(b.len())); + let x = Fp::

(rng.gen_range(0..P)); + let a = eval(&a, x); + let b = eval(&b, x); + let c = eval(&c, x); + assert_eq!(a + b, c); + } + } + + #[test] + fn test_sub() { + const P: u64 = 13; + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..1000 { + let n = rng.gen_range(0..10); + let a = gen_poly::

(n, &mut rng); + let b = gen_poly::

(n, &mut rng); + let c = sub(&a, &b); + assert!(c.len() <= a.len().max(b.len())); + let x = Fp::

(rng.gen_range(0..P)); + let a = eval(&a, x); + let b = eval(&b, x); + let c = eval(&c, x); + assert_eq!(a - b, c); + } + } + + #[test] + fn test_mul() { + const P: u64 = 13; + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..1000 { + let n = rng.gen_range(0..10); + let m = rng.gen_range(0..10); + let a = gen_poly::

(n, &mut rng); + let b = gen_poly::

(m, &mut rng); + let c = mul(&a, &b); + if a.is_empty() || b.is_empty() { + assert_eq!(c.len(), 0); + } else { + assert_eq!(c.len(), a.len() + b.len() - 1); + } + let x = Fp::

(rng.gen_range(0..P)); + let a = eval(&a, x); + let b = eval(&b, x); + let c = eval(&c, x); + assert_eq!(a * b, c); + } + } + + #[test] + fn test_div() { + const P: u64 = 13; + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..1000 { + let n = rng.gen_range(0..10); + let m = rng.gen_range(1..10); + let a = gen_poly::

(n, &mut rng); + let b = gen_poly::

(m, &mut rng); + let mut r = a.clone(); + let c = div(&mut r, &b); + assert_eq!(c.len(), a.len().saturating_sub(b.len() - 1)); + let x = Fp::

(rng.gen_range(0..P)); + let a = eval(&a, x); + let b = eval(&b, x); + let c = eval(&c, x); + let r = eval(&r, x); + assert_eq!(a, b * c + r); + } + } + + #[derive(Debug, Clone, Copy, PartialEq)] + struct MinPlus(u64); + impl Add for MinPlus { + type Output = Self; + + fn add(self, rhs: Self) -> Self { + MinPlus(self.0.min(rhs.0)) + } + } + impl AddAssign for MinPlus { + fn add_assign(&mut self, rhs: Self) { + *self = *self + rhs; + } + } + impl Sum for MinPlus { + fn sum>(iter: I) -> Self { + iter.fold(MinPlus(std::u64::MAX), |a, b| a + b) + } + } + impl Mul for MinPlus { + type Output = Self; + + fn mul(self, rhs: Self) -> Self { + MinPlus(self.0.saturating_add(rhs.0)) + } + } + impl MulAssign for MinPlus { + fn mul_assign(&mut self, rhs: Self) { + *self = *self * rhs; + } + } + impl Product for MinPlus { + fn product>(iter: I) -> Self { + iter.fold(MinPlus(0), |a, b| a * b) + } + } + + #[test] + fn test_add_min_plus() { + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..1000 { + let n = rng.gen_range(0..10); + let a = (0..n) + .map(|_| MinPlus(rng.gen_range(0..100))) + .collect::>(); + let b = (0..n) + .map(|_| MinPlus(rng.gen_range(0..100))) + .collect::>(); + let c = add(&a, &b); + assert_eq!(c.len(), a.len().max(b.len())); + for i in 0..c.len() { + assert_eq!(c[i].0, a[i].0.min(b[i].0),); + } + } + } + + #[test] + fn test_mul_min_plus() { + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..1000 { + let n = rng.gen_range(0..10); + let m = rng.gen_range(0..10); + let a = (0..n) + .map(|_| MinPlus(rng.gen_range(0..100))) + .collect::>(); + let b = (0..m) + .map(|_| MinPlus(rng.gen_range(0..100))) + .collect::>(); + let c = mul(&a, &b); + if a.is_empty() || b.is_empty() { + assert_eq!(c.len(), 0); + } else { + assert_eq!(c.len(), a.len() + b.len() - 1); + } + for k in 0..c.len() { + let mut x = std::u64::MAX; + for i in k.saturating_sub(b.len() - 1)..=k.min(a.len() - 1) { + x = x.min(a[i].0 + b[k - i].0); + } + assert_eq!(c[k].0, x); + } + } + } + + #[test] + fn test_pow_min_plus() { + let mut rng = StdRng::seed_from_u64(42); + for _ in 0..1000 { + let n = rng.gen_range(0..10); + let a = (0..n) + .map(|_| MinPlus(rng.gen_range(0..100))) + .collect::>(); + let e = rng.gen_range(0..10); + let b = pow(a.clone(), e); + if a.is_empty() { + vec![MinPlus(0)]; + } else { + let mut c = vec![MinPlus(u64::MAX); (a.len() - 1) * e as usize + 1]; + c[0] = MinPlus(0); + for _ in 0..e { + let mut swp = vec![MinPlus(u64::MAX); c.len()]; + for i in 0..=c.len() - a.len() { + for j in 0..a.len() { + swp[i + j].0 = swp[i + j].0.min(c[i].0.saturating_add(a[j].0)); + } + } + c = swp; + } + assert_eq!(b, c); + } + } + } +} diff --git a/libs/poly/src/lib.rs b/libs/poly/src/lib.rs deleted file mode 100644 index c9e34f7f..00000000 --- a/libs/poly/src/lib.rs +++ /dev/null @@ -1,303 +0,0 @@ -//! [`std::ops`] を利用した多項式のナイーブな計算 -//! -//! # あるもの -//! -//! -//! | お名前 | 計算量 | やること | -//! | - | - | - | -//! | [`poly_add`] | Θ ( max ( l, m ) ) | 足し算 | -//! | [`poly_sub`] | Θ ( max ( l, m ) ) | 引き算 | -//! | [`poly_mul`] | Θ ( l + m ) | 掛け算 | - -use std::fmt::Debug; -use std::iter::repeat_with; -use std::iter::Fuse; -use std::marker::PhantomData; -use std::ops::Add; -use std::ops::AddAssign; -use std::ops::Mul; -use std::ops::Neg; -use std::ops::Sub; - -/// 自由関数 [`poly_add`], [`poly_sub`] を呼び出します。 -/// -/// # Examples -/// -/// ``` -/// use poly::Poly; -/// -/// // 成分ごとの和を計算します。 -/// let a = vec![1, 2]; -/// let b = vec![10]; -/// let c = a.into_iter().poly_add(b).collect::>(); -/// assert_eq!(c, vec![11, 2]); -/// -/// let a = vec![1, 2]; -/// let b = vec![10]; -/// let c = a.into_iter().poly_sub(b).collect::>(); -/// assert_eq!(c, vec![-9, 2]); -/// ``` -pub trait Poly: Iterator + Sized { - fn poly_add( - self, - b: impl IntoIterator, - ) -> PolyAdd - where - Self::Item: Add, - B: Iterator, - { - poly_add(self, b) - } - fn poly_sub( - self, - b: impl IntoIterator, - ) -> PolySub - where - Self::Item: Sub + Neg, - B: Iterator, - { - poly_sub(self, b) - } -} -impl Poly for T {} - -/// 足し算 -/// -/// 成分ごとの和（[`Add::add`] -/// の結果を）を計算します。長さが異なる場合は、片方が都議らたあとはもう一方の要素要素をそのまま返します。 -/// -/// # Examples -/// -/// ``` -/// # use poly::poly_add; -/// // 成分ごとの和を計算します。 -/// let a = vec![1, 2]; -/// let b = vec![10]; -/// let c = poly_add(a, b).collect::>(); -/// assert_eq!(c, vec![11, 2]); -/// ``` -pub fn poly_add( - a: impl IntoIterator, - b: impl IntoIterator, -) -> PolyAdd -where - T: Add, - A: Iterator, - B: Iterator, -{ - PolyAdd { - a: a.into_iter().fuse(), - b: b.into_iter().fuse(), - __marker: PhantomData, - } -} -/// [`poly_add` の戻り値型ですので、そちらのドキュメントをご覧ください。](poly_add) -pub struct PolyAdd { - a: Fuse, - b: Fuse, - __marker: PhantomData, -} -impl Iterator for PolyAdd -where - T: Add, - A: Iterator, - B: Iterator, -{ - type Item = T; - - fn next(&mut self) -> Option { - match (self.a.next(), self.b.next()) { - (None, None) => None, - (Some(a), None) => Some(a), - (None, Some(b)) => Some(b), - (Some(a), Some(b)) => Some(a + b), - } - } - - fn size_hint(&self) -> (usize, Option) { - let (a_lower, a_upper) = self.a.size_hint(); - let (b_lower, b_upper) = self.b.size_hint(); - let lower = std::cmp::max(a_lower, b_lower); - let upper = match (a_upper, b_upper) { - (Some(x), Some(y)) => Some(std::cmp::max(x, y)), - _ => None, - }; - (lower, upper) - } -} - -/// 引き算 -/// -/// 成分ごとの差（[`Sub::sub`] -/// の結果を）を計算します。長さが異なる場合の仕様は [`poly_sub`] に似ていますが、`a` -/// が余ればそのまま、`b` が余ればその negation（[`Neg::neg`] の結果）を返します。 -/// -/// # Examples -/// -/// ``` -/// # use poly::poly_sub; -/// // 成分ごとの差を計算します。 -/// let a = vec![1, 2]; -/// let b = vec![10]; -/// let c = poly_sub(a, b).collect::>(); -/// assert_eq!(c, vec![-9, 2]); -/// ``` -pub fn poly_sub( - a: impl IntoIterator, - b: impl IntoIterator, -) -> PolySub -where - T: Sub, - A: Iterator, - B: Iterator, -{ - PolySub { - a: a.into_iter().fuse(), - b: b.into_iter().fuse(), - __marker: PhantomData, - } -} -/// [`poly_sub` の戻り値型ですので、そちらのドキュメントをご覧ください。](poly_sub) -pub struct PolySub { - a: Fuse, - b: Fuse, - __marker: PhantomData, -} -impl Iterator for PolySub -where - T: Sub + Neg, - A: Iterator, - B: Iterator, -{ - type Item = T; - - fn next(&mut self) -> Option { - match (self.a.next(), self.b.next()) { - (None, None) => None, - (Some(a), None) => Some(a), - (None, Some(b)) => Some(-b), - (Some(a), Some(b)) => Some(a - b), - } - } - - fn size_hint(&self) -> (usize, Option) { - let (a_lower, a_upper) = self.a.size_hint(); - let (b_lower, b_upper) = self.b.size_hint(); - let lower = std::cmp::max(a_lower, b_lower); - let upper = match (a_upper, b_upper) { - (Some(x), Some(y)) => Some(std::cmp::max(x, y)), - _ => None, - }; - (lower, upper) - } -} - -/// 掛け算 -/// -/// 多項式としての積を計算します。長さは、`a` と `b` のいずれかが殻ならば空、さもなくば `a.len() + -/// b.len() - 1` となります。 -/// -/// # Examples -/// -/// ``` -/// # use poly::poly_mul; -/// // 多項式の積を計算します。 -/// let a = vec![1, 2]; -/// let b = vec![10, 20]; -/// let c = poly_mul(&a, &b, || 0); -/// assert_eq!(c, vec![10, 40, 40]); -/// ``` -pub fn poly_mul + Debug + Copy>( - a: &[T], - b: &[T], - zero: impl Fn() -> T, -) -> Vec { - if a.is_empty() || b.is_empty() { - Vec::new() - } else { - let mut c = repeat_with(zero) - .take(a.len() + b.len() - 1) - .collect::>(); - a.iter() - .enumerate() - .for_each(|(i, &x)| c[i..].iter_mut().zip(b).for_each(|(z, &y)| *z += x * y)); - c - } -} - -#[cfg(test)] -mod tests { - use super::poly_add; - use super::poly_mul; - use super::poly_sub; - use rand::prelude::StdRng; - use rand::Rng; - use rand::SeedableRng; - use std::iter::repeat_with; - - #[test] - fn test_add() { - let mut rng = StdRng::seed_from_u64(42); - for _ in 0..20 { - let l = rng.gen_range(1..10); - let m = rng.gen_range(1..10); - let a = repeat_with(|| rng.gen_range(-10..=10)) - .take(l) - .collect::>(); - let b = repeat_with(|| rng.gen_range(-10..=10)) - .take(m) - .collect::>(); - let mut c = vec![0; a.len().max(b.len())]; - c.iter_mut().zip(&a).for_each(|(z, &x)| *z += x); - c.iter_mut().zip(&b).for_each(|(z, &y)| *z += y); - let expected = c; - let result = poly_add(a, b).collect::>(); - assert_eq!(result, expected); - } - } - - #[test] - fn test_sub() { - let mut rng = StdRng::seed_from_u64(42); - for _ in 0..20 { - let l = rng.gen_range(1..10); - let m = rng.gen_range(1..10); - let a = repeat_with(|| rng.gen_range(-10..=10)) - .take(l) - .collect::>(); - let b = repeat_with(|| rng.gen_range(-10..=10)) - .take(m) - .collect::>(); - let mut c = vec![0; a.len().max(b.len())]; - c.iter_mut().zip(&a).for_each(|(z, &x)| *z += x); - c.iter_mut().zip(&b).for_each(|(z, &y)| *z -= y); - let expected = c; - let result = poly_sub(a, b).collect::>(); - assert_eq!(result, expected); - } - } - - #[test] - fn test_mul() { - let mut rng = StdRng::seed_from_u64(42); - for _ in 0..20 { - let l = rng.gen_range(1..10); - let m = rng.gen_range(1..10); - let a = repeat_with(|| rng.gen_range(-10..=10)) - .take(l) - .collect::>(); - let b = repeat_with(|| rng.gen_range(-10..=10)) - .take(m) - .collect::>(); - let mut c = vec![0; a.len() + b.len() - 1]; - for i in 0..l { - for j in 0..m { - c[i + j] += a[i] * b[j]; - } - } - let expected = c; - let result = poly_mul(&a, &b, || 0); - assert_eq!(result, expected); - } - } -}