結果
問題 | No.754 畳み込みの和 |
ユーザー | ngtkana |
提出日時 | 2024-05-06 19:57:36 |
言語 | Rust (1.77.0 + proconio) |
結果 |
AC
|
実行時間 | 2,730 ms / 5,000 ms |
コード長 | 18,476 bytes |
コンパイル時間 | 12,296 ms |
コンパイル使用メモリ | 401,388 KB |
実行使用メモリ | 13,124 KB |
最終ジャッジ日時 | 2024-05-06 19:58:14 |
合計ジャッジ時間 | 23,902 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge1 |
(要ログイン)
テストケース
テストケース表示入力 | 結果 | 実行時間 実行使用メモリ |
---|---|---|
testcase_00 | AC | 2,636 ms
13,120 KB |
testcase_01 | AC | 2,730 ms
13,084 KB |
testcase_02 | AC | 2,609 ms
13,124 KB |
コンパイルメッセージ
warning: unused import: `factorial::Factorial` --> src/main.rs:325:13 | 325 | pub use factorial::Factorial; | ^^^^^^^^^^^^^^^^^^^^ | = note: `#[warn(unused_imports)]` on by default warning: unused import: `fourier::any_mod_fps_mul` --> src/main.rs:326:13 | 326 | pub use fourier::any_mod_fps_mul; | ^^^^^^^^^^^^^^^^^^^^^^^^ warning: unused import: `fourier::fft` --> src/main.rs:327:13 | 327 | pub use fourier::fft; | ^^^^^^^^^^^^ warning: unused import: `fourier::fps_mul` --> src/main.rs:328:13 | 328 | pub use fourier::fps_mul; | ^^^^^^^^^^^^^^^^ warning: unused import: `fourier::ifft` --> src/main.rs:329:13 | 329 | pub use fourier::ifft; | ^^^^^^^^^^^^^
ソースコード
use proconio::input; use std::ops::Add; use std::ops::Mul; use std::ops::Sub; type Fp = fp::Fp<1000000007>; fn main() { input! { n: usize, a: [u64; n + 1], b: [u64; n + 1], } let a = a.iter().map(|&x| Fp::new(x)).collect::<Vec<_>>(); let b = b.iter().map(|&x| Fp::new(x)).collect::<Vec<_>>(); let c = karatsuba(&a, &b); let ans = c[..=n].iter().sum::<Fp>(); println!("{}", ans); } trait Zero { fn zero() -> Self; } impl Zero for Fp { fn zero() -> Self { Fp::new(0) } } fn karatsuba<T>(a: &[T], b: &[T]) -> Vec<T> where T: Clone + Copy + Mul<Output = T> + Add<Output = T> + Sub<Output = T> + Zero + std::fmt::Debug, { let n = a.len().max(b.len()).next_power_of_two(); let mut slice = vec![T::zero(); 4 * n - 2]; for (i, &x) in a.iter().enumerate() { slice[2 * (n - 1 + i)] = x; } for (i, &x) in b.iter().enumerate() { slice[2 * (n - 1 + i) + 1] = x; } let mut left = 0; let mut right = 2 * (n - 1); let mut structure = 0; let mut size = n; while structure != n * n { while size != 1 { size >>= 1; right -= 2 * size; for ((i, j), k) in (right..) .zip(right + 2 * size..) .zip(right + 4 * size..) .take(2 * size) { slice[i] = slice[j]; slice[j] = slice[k]; slice[k] = slice[i] + slice[j]; } } let (x, y) = (slice[left], slice[left + 1]); (slice[left], slice[left + 1]) = (x * y, T::zero()); left += 2; right += 2; structure += 1; while structure & (3 * size * size) == 3 * size * size { for ((i, j), k) in (left - 6 * size..) .zip(left - 4 * size..) .zip(left - 2 * size..) .take(2 * size) { slice[k] = slice[k] - slice[i] - slice[j]; } for (i, j) in (left - 5 * size..).zip(left - 2 * size..).take(2 * size) { slice[i] = slice[i] + slice[j]; slice[j] = T::zero(); } left -= 2 * size; structure += size * size; size <<= 1; } } slice[..2 * n - 1].to_vec() } // fp {{{ // https://ngtkana.github.io/ac-adapter-rs/fp/index.html #[allow(dead_code)] mod fp { mod ext_gcd { pub(crate) fn mod_inv<const P: u64>(x: u64) -> u64 { debug_assert!(P % 2 == 1); debug_assert!(P < 1 << 31); debug_assert!(x < P); mod_inv_signed(x as i64, P as i64) as u64 } fn mod_inv_signed(a: i64, m: i64) -> i64 { debug_assert!(a > 0); debug_assert!(m > 0); if a == 1 { return 1; } m + (1 - m * mod_inv_signed(m % a, a)) / a } } mod factorial { use super::Fp; use std::ops::Index; pub struct Factorial<const P: u64> { fact: Vec<Fp<P>>, inv_fact: Vec<Fp<P>>, } impl<const P: u64> Factorial<P> { pub fn new(length: usize) -> Self { let mut fact = vec![Fp::<P>::new(1); length + 1]; let mut inv_fact = vec![Fp::<P>::new(1); length + 1]; for i in 1..=length { fact[i] = fact[i - 1] * Fp::<P>::new(i as u64); } inv_fact[length] = fact[length].inv(); for i in (1..=length).rev() { inv_fact[i - 1] = inv_fact[i] * Fp::<P>::new(i as u64); } Self { fact, inv_fact } } pub fn fact(&self, n: usize) -> Fp<P> { self.fact[n] } pub fn inv_fact(&self, n: usize) -> Fp<P> { self.inv_fact[n] } pub fn perm(&self, n: usize, k: usize) -> Fp<P> { self.fact[n] * self.inv_fact[n - k] } pub fn comb(&self, n: usize, k: usize) -> Fp<P> { self.fact[n] * self.inv_fact[n - k] * self.inv_fact[k] } pub fn binom(&self, n: usize, k: usize) -> Fp<P> { self.comb(n, k) } pub fn comb_or_zero(&self, n: usize, k: isize) -> Fp<P> { if k < 0 || k as usize > n { Fp::<P>::new(0) } else { self.comb(n, k as usize) } } pub fn comb_with_reputation(&self, n: usize, k: usize) -> Fp<P> { assert!(n > 0 || k > 0); self.comb(n + k - 1, k) } } impl<const P: u64> Index<usize> for Factorial<P> { type Output = Fp<P>; fn index(&self, index: usize) -> &Self::Output { &self.fact[index] } } } mod fourier { use super::mod_inv; use super::Fp; use super::PrimitiveRoot; const P1: u64 = 924844033; const P2: u64 = 998244353; const P3: u64 = 1012924417; type F1 = Fp<P1>; type F2 = Fp<P2>; type F3 = Fp<P3>; pub fn fps_mul<const P: u64>(a: impl AsRef<[Fp<P>]>, b: impl AsRef<[Fp<P>]>) -> Vec<Fp<P>> where (): PrimitiveRoot<P>, { let a = a.as_ref(); let b = b.as_ref(); if a.is_empty() || b.is_empty() { return vec![]; } let mut a = a.to_vec(); let mut b = b.to_vec(); let n = a.len() + b.len() - 1; let len = n.next_power_of_two(); a.resize(len, Fp::new(0)); b.resize(len, Fp::new(0)); fft(&mut a); fft(&mut b); for (a, b) in a.iter_mut().zip(b.iter()) { *a *= *b; } ifft(&mut a); a.truncate(n); a } pub fn any_mod_fps_mul<const P: u64>(a: &[Fp<P>], b: &[Fp<P>]) -> Vec<Fp<P>> { let v1 = fps_mul( a.iter().map(|&x| F1::new(x.value())).collect::<Vec<_>>(), b.iter().map(|&x| F1::new(x.value())).collect::<Vec<_>>(), ); let v2 = fps_mul( a.iter().map(|&x| F2::new(x.value())).collect::<Vec<_>>(), b.iter().map(|&x| F2::new(x.value())).collect::<Vec<_>>(), ); let v3 = fps_mul( a.iter().map(|&x| F3::new(x.value())).collect::<Vec<_>>(), b.iter().map(|&x| F3::new(x.value())).collect::<Vec<_>>(), ); v1.into_iter() .zip(v2) .zip(v3) .map(|((e1, e2), e3)| garner(e1, e2, e3)) .collect::<Vec<_>>() } pub fn fft<const P: u64>(f: &mut [Fp<P>]) where (): PrimitiveRoot<P>, { let n = f.len(); assert!(n.is_power_of_two()); assert!((P - 1) % n as u64 == 0); let mut root = <() as PrimitiveRoot<P>>::VALUE.pow((P - 1) / f.len() as u64); let fourth = <() as PrimitiveRoot<P>>::VALUE.pow((P - 1) / 4); let mut fft_len = n; while 4 <= fft_len { let quarter = fft_len / 4; for f in f.chunks_mut(fft_len) { let mut c = Fp::new(1); for (((i, j), k), l) in (0..) .zip(quarter..) .zip(quarter * 2..) .zip(quarter * 3..) .take(quarter) { let c2 = c * c; let x = f[i] + f[k]; let y = f[j] + f[l]; let z = f[i] - f[k]; let w = fourth * (f[j] - f[l]); f[i] = x + y; f[j] = c2 * (x - y); f[k] = c * (z + w); f[l] = c2 * c * (z - w); c *= root; } } root *= root; root *= root; fft_len = quarter; } if fft_len == 2 { for f in f.chunks_mut(2) { let x = f[0]; let y = f[1]; f[0] = x + y; f[1] = x - y; } } } pub fn ifft<const P: u64>(f: &mut [Fp<P>]) where (): PrimitiveRoot<P>, { let n = f.len(); assert!(n.is_power_of_two()); let root = <() as PrimitiveRoot<P>>::VALUE.pow((P - 1) / f.len() as u64); let mut roots = std::iter::successors(Some(root.inv()), |x| Some(x * x)) .take(n.trailing_zeros() as usize + 1) .collect::<Vec<_>>(); roots.reverse(); let fourth = <() as PrimitiveRoot<P>>::VALUE.pow((P - 1) / 4).inv(); let mut quarter = 1_usize; if n.trailing_zeros() % 2 == 1 { for f in f.chunks_mut(2) { let x = f[0]; let y = f[1]; f[0] = x + y; f[1] = x - y; } quarter = 2; } while quarter != n { let fft_len = quarter * 4; let root = roots[fft_len.trailing_zeros() as usize]; for f in f.chunks_mut(fft_len) { let mut c = Fp::new(1); for (((i, j), k), l) in (0..) .zip(quarter..) .zip(quarter * 2..) .zip(quarter * 3..) .take(quarter) { let c2 = c * c; let x = f[i] + c2 * f[j]; let y = f[i] - c2 * f[j]; let z = c * (f[k] + c2 * f[l]); let w = fourth * c * (f[k] - c2 * f[l]); f[i] = x + z; f[j] = y + w; f[k] = x - z; f[l] = y - w; c *= root; } } quarter = fft_len; } let d = Fp::from(f.len()).inv(); f.iter_mut().for_each(|x| *x *= d); } fn garner<const P: u64>(x1: Fp<P1>, x2: Fp<P2>, x3: Fp<P3>) -> Fp<P> { let (x1, x2, x3) = (x1.value(), x2.value(), x3.value()); let x2 = ((x2 + (P2 - x1)) * mod_inv::<P2>(P1)) % P2; let x3 = (((x3 + (P3 - x1)) * mod_inv::<P3>(P1) % P3 + (P3 - x2)) * mod_inv::<P3>(P2)) % P3; Fp::new(x1 + P1 * (x2 + P2 * x3 % P)) } } use ext_gcd::mod_inv; pub use factorial::Factorial; pub use fourier::any_mod_fps_mul; pub use fourier::fft; pub use fourier::fps_mul; pub use fourier::ifft; use std::iter::Product; use std::iter::Sum; use std::mem::swap; use std::ops::Add; use std::ops::AddAssign; use std::ops::Div; use std::ops::DivAssign; use std::ops::Mul; use std::ops::MulAssign; use std::ops::Neg; use std::ops::Sub; use std::ops::SubAssign; #[macro_export] macro_rules! fp { ($value:expr) => { $crate::fp::Fp::from($value) }; ($value:expr; mod $p:expr) => { $crate::fp::Fp::<$p>::from($value) }; } pub trait PrimitiveRoot<const P: u64> { const VALUE: Fp<P>; } impl PrimitiveRoot<998244353> for () { const VALUE: Fp<998244353> = Fp::new(3); } impl PrimitiveRoot<1012924417> for () { const VALUE: Fp<1012924417> = Fp::new(5); } impl PrimitiveRoot<924844033> for () { const VALUE: Fp<924844033> = Fp::new(5); } #[derive(Clone, Copy, PartialEq, Eq, Hash)] pub struct Fp<const P: u64> { value: u64, } impl<const P: u64> Fp<P> { pub const fn new(value: u64) -> Self { Self { value: value % P } } pub const fn value(self) -> u64 { self.value } pub fn inv(self) -> Self { Self { value: mod_inv::<P>(self.value), } } pub fn pow(self, mut exp: u64) -> Self { let mut result = Self::new(1); let mut base = self; while exp > 0 { if exp & 1 == 1 { result *= base; } base *= base; exp >>= 1; } result } pub fn sign(pow: usize) -> Self { Self::new(if pow % 2 == 0 { 1 } else { P - 1 }) } } impl<const P: u64> std::fmt::Debug for Fp<P> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { pub fn berlekamp_massey_fp(a: i64, p: i64) -> [i64; 2] { let mut u0 = 0_i64; let mut v0 = 1_i64; let mut w0 = a * u0 + p * v0; let mut u1 = 1_i64; let mut v1 = 0_i64; let mut w1 = a * u1 + p * v1; while p <= w0 * w0 { let q = w0 / w1; u0 -= q * u1; v0 -= q * v1; w0 -= q * w1; swap(&mut u0, &mut u1); swap(&mut v0, &mut v1); swap(&mut w0, &mut w1); } [w0, u0] } if self.value == 0 { return write!(f, "0"); } let [mut num, mut den] = berlekamp_massey_fp(self.value as i64, P as i64); if den < 0 { num = -num; den = -den; } if den == 1 { write!(f, "{}", num) } else { write!(f, "{}/{}", num, den) } } } impl<const P: u64> std::fmt::Display for Fp<P> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.value()) } } macro_rules! impl_from_signed { ($($t:ty),*) => { $( impl<const P: u64> From<$t> for Fp<P> { fn from(x: $t) -> Self { if x < 0 { -Self::new((P as i64 - x as i64) as u64) } else { Self::new(x as u64) } } } )* }; } impl_from_signed!(i8, i16, i32, i64, i128, isize); macro_rules! impl_from_unsigned { ($($t:ty),*) => { $( impl<const P: u64> From<$t> for Fp<P> { fn from(x: $t) -> Self { Self::new(x as u64) } } )* }; } impl_from_unsigned!(u8, u16, u32, u64, u128, usize); impl<const P: u64> AddAssign<Fp<P>> for Fp<P> { fn add_assign(&mut self, rhs: Fp<P>) { self.value += rhs.value; if self.value >= P { self.value -= P; } } } impl<const P: u64> SubAssign<Fp<P>> for Fp<P> { fn sub_assign(&mut self, rhs: Fp<P>) { if self.value < rhs.value { self.value += P; } self.value -= rhs.value; } } impl<const P: u64> MulAssign<Fp<P>> for Fp<P> { fn mul_assign(&mut self, rhs: Fp<P>) { self.value = self.value * rhs.value % P; } } #[allow(clippy::suspicious_op_assign_impl)] impl<const P: u64> DivAssign<Fp<P>> for Fp<P> { fn div_assign(&mut self, rhs: Fp<P>) { *self *= rhs.inv() } } macro_rules! fp_forward_ops { ($( $trait:ident, $trait_assign:ident, $fn:ident, $fn_assign:ident, )*) => {$( impl<const P: u64> $trait_assign<&Fp<P>> for Fp<P> { fn $fn_assign(&mut self, rhs: &Fp<P>) { self.$fn_assign(*rhs); } } impl<const P: u64, T: Into<Fp<P>>> $trait<T> for Fp<P> { type Output = Fp<P>; fn $fn(mut self, rhs: T) -> Self::Output { self.$fn_assign(rhs.into()); self } } impl<const P: u64> $trait<&Fp<P>> for Fp<P> { type Output = Fp<P>; fn $fn(self, rhs: &Fp<P>) -> Self::Output { self.$fn(*rhs) } } impl<const P: u64, T: Into<Fp<P>>> $trait<T> for &Fp<P> { type Output = Fp<P>; fn $fn(self, rhs: T) -> Self::Output { (*self).$fn(rhs.into()) } } impl<const P: u64> $trait<&Fp<P>> for &Fp<P> { type Output = Fp<P>; fn $fn(self, rhs: &Fp<P>) -> Self::Output { (*self).$fn(*rhs) } } )*}; } fp_forward_ops! { Add, AddAssign, add, add_assign, Sub, SubAssign, sub, sub_assign, Mul, MulAssign, mul, mul_assign, Div, DivAssign, div, div_assign, } impl<const P: u64> Neg for Fp<P> { type Output = Fp<P>; fn neg(mut self) -> Self::Output { if self.value > 0 { self.value = P - self.value; } self } } impl<const P: u64> Sum for Fp<P> { fn sum<I: Iterator<Item = Self>>(iter: I) -> Self { iter.fold(Self::new(0), |acc, x| acc + x) } } impl<'a, const P: u64> Sum<&'a Self> for Fp<P> { fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self { iter.copied().sum() } } impl<const P: u64> Product for Fp<P> { fn product<I: Iterator<Item = Self>>(iter: I) -> Self { iter.fold(Self::new(1), |acc, x| acc * x) } } impl<'a, const P: u64> Product<&'a Self> for Fp<P> { fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self { iter.copied().product() } } } // }}}