結果

問題 No.840 ほむほむほむら
ユーザー ngtkanangtkana
提出日時 2024-05-24 03:01:29
言語 Rust
(1.77.0 + proconio)
結果
AC  
実行時間 107 ms / 4,000 ms
コード長 20,194 bytes
コンパイル時間 14,253 ms
コンパイル使用メモリ 402,892 KB
実行使用メモリ 6,948 KB
最終ジャッジ日時 2024-05-24 03:01:46
合計ジャッジ時間 16,342 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 12 ms
6,816 KB
testcase_01 AC 12 ms
6,812 KB
testcase_02 AC 13 ms
6,944 KB
testcase_03 AC 12 ms
6,944 KB
testcase_04 AC 12 ms
6,944 KB
testcase_05 AC 23 ms
6,944 KB
testcase_06 AC 30 ms
6,940 KB
testcase_07 AC 31 ms
6,940 KB
testcase_08 AC 24 ms
6,940 KB
testcase_09 AC 30 ms
6,940 KB
testcase_10 AC 59 ms
6,944 KB
testcase_11 AC 60 ms
6,944 KB
testcase_12 AC 56 ms
6,940 KB
testcase_13 AC 63 ms
6,940 KB
testcase_14 AC 63 ms
6,940 KB
testcase_15 AC 95 ms
6,944 KB
testcase_16 AC 98 ms
6,940 KB
testcase_17 AC 106 ms
6,940 KB
testcase_18 AC 106 ms
6,940 KB
testcase_19 AC 106 ms
6,944 KB
testcase_20 AC 2 ms
6,940 KB
testcase_21 AC 5 ms
6,940 KB
testcase_22 AC 5 ms
6,948 KB
testcase_23 AC 107 ms
6,940 KB
testcase_24 AC 107 ms
6,944 KB
testcase_25 AC 8 ms
6,944 KB
testcase_26 AC 8 ms
6,944 KB
testcase_27 AC 89 ms
6,940 KB
権限があれば一括ダウンロードができます
コンパイルメッセージ
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;
    |             ^^^^^^^^^^^^^

ソースコード

diff # 

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()
        }
    }
}
// }}}
0