結果
問題 | No.840 ほむほむほむら |
ユーザー |
![]() |
提出日時 | 2024-05-24 03:01:29 |
言語 | Rust (1.83.0 + proconio) |
結果 |
AC
|
実行時間 | 111 ms / 4,000 ms |
コード長 | 20,194 bytes |
コンパイル時間 | 13,765 ms |
コンパイル使用メモリ | 384,748 KB |
実行使用メモリ | 5,248 KB |
最終ジャッジ日時 | 2024-12-20 19:12:51 |
合計ジャッジ時間 | 16,365 ms |
ジャッジサーバーID (参考情報) |
judge3 / judge5 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 3 |
other | AC * 25 |
コンパイルメッセージ
warning: unused import: `factorial::Factorial` --> src/main.rs:385:13 | 385 | pub use factorial::Factorial; | ^^^^^^^^^^^^^^^^^^^^ | = note: `#[warn(unused_imports)]` on by default warning: unused import: `fourier::any_mod_fps_mul` --> src/main.rs:386:13 | 386 | pub use fourier::any_mod_fps_mul; | ^^^^^^^^^^^^^^^^^^^^^^^^ warning: unused import: `fourier::fft` --> src/main.rs:387:13 | 387 | pub use fourier::fft; | ^^^^^^^^^^^^ warning: unused import: `fourier::fps_mul` --> src/main.rs:388:13 | 388 | pub use fourier::fps_mul; | ^^^^^^^^^^^^^^^^ warning: unused import: `fourier::ifft` --> src/main.rs:389:13 | 389 | pub use fourier::ifft; | ^^^^^^^^^^^^^
ソースコード
type Fp = fp::Fp<998244353>;const D: usize = 5;const D3: usize = D * D * D;fn main() {let (e, d): (usize, usize) = io::input();let encode = |i: usize, j: usize, k: usize| i * d * d + j * d + k;let mut a = [[Fp::new(0); D3]; D3];for i in 0..d {for j in 0..d {for k in 0..d {let from = encode(i, j, k);for to in [encode((i + 1) % d, j, k),encode(i, (j + i) % d, k),encode(i, j, (k + j) % d),] {a[from][to] += Fp::new(1);}}}}let mut x = [Fp::new(0); D3];x[encode(0, 0, 0)] = Fp::new(1);let x = pow_apply(a, x, e);let ans = (0..d).map(|i| (0..d).map(|j| x[encode(i, j, 0)]).sum::<Fp>()).sum::<Fp>();println!("{}", ans);}fn pow_apply(mut a: [[Fp; D3]; D3], mut x: [Fp; D3], mut n: usize) -> [Fp; D3] {assert_ne!(n, 0);while n > 1 {if n & 1 == 1 {x = apply(&a, &x);}a = mul(&a, &a);n >>= 1;}apply(&a, &x)}fn apply(a: &[[Fp; D3]; D3], x: &[Fp; D3]) -> [Fp; D3] {let mut y = [Fp::new(0); D3];for i in 0..D3 {for j in 0..D3 {y[i] += a[i][j] * x[j];}}y}fn mul(a: &[[Fp; D3]; D3], b: &[[Fp; D3]; D3]) -> [[Fp; D3]; D3] {let n = a.len();let mut c = [[Fp::new(0); D3]; D3];for i in 0..n {for j in 0..n {for k in 0..n {c[i][j] += a[i][k] * b[k][j];}}}c}// io {{{// https://ngtkana.github.io/ac-adapter-rs/io/index.html#[allow(dead_code)]mod io {use std::cell::Cell;use std::io::stdin;use std::io::BufRead;use std::io::BufReader;use std::io::Lines;use std::io::Stdin;use std::sync::Mutex;use std::sync::Once;pub fn input<T: ParseLine>() -> T {ParseLine::parse_line(&line())}pub trait ParseLine {fn parse_line(s: &str) -> Self;}macro_rules! impl_parse_line {($($t:ty),*) => {$(impl ParseLine for $t {fn parse_line(s: &str) -> Self {s.parse().unwrap()}})*};}impl_parse_line!(u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize, String, char);macro_rules! impl_parse_line_tuple {($($t:ident),*) => {impl<$($t: ParseLine),*> ParseLine for ($($t,)*) {fn parse_line(s: &str) -> Self {let mut s = s.split_whitespace();($($t::parse_line(s.next().unwrap()),)*)}}};}impl_parse_line_tuple!(T0, T1);impl_parse_line_tuple!(T0, T1, T2);impl_parse_line_tuple!(T0, T1, T2, T3);impl_parse_line_tuple!(T0, T1, T2, T3, T4);impl_parse_line_tuple!(T0, T1, T2, T3, T4, T5);impl_parse_line_tuple!(T0, T1, T2, T3, T4, T5, T6);impl_parse_line_tuple!(T0, T1, T2, T3, T4, T5, T6, T7);impl_parse_line_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8);impl_parse_line_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9);impl<T: ParseLine> ParseLine for Vec<T> {fn parse_line(s: &str) -> Self {s.split_whitespace().map(T::parse_line).collect()}}static ONCE: Once = Once::new();type Server = Mutex<Lines<BufReader<Stdin>>>;struct Lazy(Cell<Option<Server>>);unsafe impl Sync for Lazy {}fn line() -> String {static SYNCER: Lazy = Lazy(Cell::new(None));ONCE.call_once(|| {SYNCER.0.set(Some(Mutex::new(BufReader::new(stdin()).lines())));});unsafe {(*SYNCER.0.as_ptr()).as_ref().unwrap().lock().unwrap().next().unwrap().unwrap()}}}// }}}// 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()}}}// }}}