結果

問題 No.3044 よくあるカエルさん
ユーザー Blue_S
提出日時 2025-01-16 06:10:18
言語 Rust
(1.83.0 + proconio)
結果
WA  
(最新)
AC  
(最初)
実行時間 -
コード長 7,638 bytes
コンパイル時間 12,428 ms
コンパイル使用メモリ 380,820 KB
実行使用メモリ 5,248 KB
最終ジャッジ日時 2025-01-18 00:35:20
合計ジャッジ時間 14,303 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 2 WA * 18
権限があれば一括ダウンロードができます

ソースコード

diff # 

use galois_field::GaloisField;
use proconio::{input, marker::Usize1};

const MOD: u64 = 998_244_353;
fn main() {
    input! {
        n: Usize1,
        t: usize,
        k: u64,
        l: u64,
    }

    let mut dp = vec![gf!(1)];
    for i in 1..t {
        let mut next = gf!(0);
        if i >= 1 {
            next += gf!(k - 1) / gf!(6) * dp[i - 1];
        }
        if i >= 2 {
            next += gf!(l - k) / gf!(6) * dp[i - 2];
        }
        if i >= t {
            next += gf!(7 - k) / gf!(6) * dp[i - t];
        }
        dp.push(next);
    }
    assert_eq!(dp.len(), t);

    if n <= t {
        println!("{}", dp[n].value());
        return;
    }
    let mut x = vec![vec![gf!(0); t]; t];

    x[0][0] = gf!(k - 1) / gf!(6);
    x[0][1] = gf!(l - k) / gf!(6);
    x[0][t - 1] = gf!(7 - k) / gf!(6);
    for i in 1..t {
        x[i][i - 1] = gf!(1);
    }

    let coef = pow(&x, n);

    let ans = coef
        .last()
        .unwrap()
        .iter()
        .zip(dp.iter().rev())
        .map(|(coef, dp)| coef * dp)
        .sum::<GaloisField<MOD>>();
    println!("{}", ans.value());
}

fn multiple(
    a: &Vec<Vec<GaloisField<MOD>>>,
    b: &Vec<Vec<GaloisField<MOD>>>,
) -> Vec<Vec<GaloisField<MOD>>> {
    let n = a.len();
    let mut res = vec![vec![gf!(0); n]; n];
    for i in 0..n {
        for j in 0..n {
            for k in 0..n {
                res[i][j] += a[i][k] * b[k][j];
            }
        }
    }
    res
}

fn pow(a: &Vec<Vec<GaloisField<MOD>>>, mut n: usize) -> Vec<Vec<GaloisField<MOD>>> {
    let m = a.len();
    let mut res = vec![vec![gf!(0); m]; m];
    for i in 0..m {
        res[i][i] = gf!(1);
    }
    let mut base = a.clone();
    while n > 0 {
        if n & 1 == 1 {
            res = multiple(&res, &base);
        }
        base = multiple(&base, &base);
        n /= 2;
    }
    res
}

pub mod galois_field {
    use std::{
        iter::{Product, Sum},
        ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
    };

    #[derive(Clone, Copy, PartialEq, Eq, Hash)]
    pub struct GaloisField<const MOD: u64> {
        value: u64,
    }

    impl<const MOD: u64> GaloisField<MOD> {
        pub fn new(value: u64) -> Self {
            Self { value: value % MOD }
        }

        pub fn value(&self) -> u64 {
            self.value
        }

        pub fn pow(&self, mut exp: u64) -> Self {
            let mut res = Self::new(1);
            let mut base = self.clone();
            while exp > 0 {
                if exp & 1 == 1 {
                    res *= base;
                }
                base *= base;
                exp >>= 1;
            }
            res
        }

        pub fn inv(&self) -> Self {
            self.pow(MOD - 2)
        }
    }

    macro_rules! impl_from_signed {
        ($($t:ty),*) => {
            $(
                impl<const MOD: u64> From<$t> for GaloisField<MOD> {
                    fn from(x: $t) -> Self {
                        if x < 0 {
                            - Self::new((MOD as i64 - x as i64) as u64)
                        } else {
                            Self::new(x as u64)
                        }
                    }
                }
            )*
        };
    }
    macro_rules! impl_from_unsigned {
        ($($t:ty),*) => {
            $(
                impl<const MOD: u64> From<$t> for GaloisField<MOD> {
                    fn from(x: $t) -> Self {
                        Self::new(x as u64)
                    }
                }
            )*
        };
    }
    impl_from_signed!(i8, i16, i32, i64, i128, isize);
    impl_from_unsigned!(u8, u16, u32, u64, u128, usize);
    #[macro_export]
    macro_rules! gf {
        ($value:expr) => {
            $crate::GaloisField::from($value)
        };
        ($value:expr; mod $p:expr) => {
            $crate::GaloisField::<$p>::from($value)
        };
    }

    impl<const MOD: u64> AddAssign<GaloisField<MOD>> for GaloisField<MOD> {
        fn add_assign(&mut self, rhs: GaloisField<MOD>) {
            self.value += rhs.value;
            if self.value >= MOD {
                self.value -= MOD;
            }
        }
    }
    impl<const MOD: u64> SubAssign<GaloisField<MOD>> for GaloisField<MOD> {
        fn sub_assign(&mut self, rhs: GaloisField<MOD>) {
            if self.value < rhs.value {
                self.value += MOD;
            }
            self.value -= rhs.value;
        }
    }
    impl<const MOD: u64> MulAssign<GaloisField<MOD>> for GaloisField<MOD> {
        fn mul_assign(&mut self, rhs: GaloisField<MOD>) {
            self.value *= rhs.value;
            self.value %= MOD;
        }
    }
    impl<const MOD: u64> DivAssign<GaloisField<MOD>> for GaloisField<MOD> {
        fn div_assign(&mut self, rhs: GaloisField<MOD>) {
            self.value *= rhs.inv().value;
            self.value %= MOD;
        }
    }
    macro_rules! gf_forward_ops {
        ($(
                $trait:ident,
                $trait_assign:ident,
                $fn:ident,
                $fn_assign:ident,
        )*) => {$(
            impl<const MOD: u64> $trait_assign<&GaloisField<MOD>> for GaloisField<MOD> {
                fn $fn_assign(&mut self, rhs: &GaloisField<MOD>) {
                    self.$fn_assign(*rhs);
                }
            }
            impl<const MOD: u64, T: Into<GaloisField<MOD>>> $trait<T> for GaloisField<MOD> {
                type Output = GaloisField<MOD>;
                fn $fn(mut self, rhs: T) -> Self::Output {
                    self.$fn_assign(rhs.into());
                    self
                }
            }
            impl<const MOD: u64> $trait<&GaloisField<MOD>> for GaloisField<MOD> {
                type Output = GaloisField<MOD>;
                fn $fn(self, rhs: &GaloisField<MOD>) -> Self::Output {
                    self.$fn(*rhs)
                }
            }
            impl<const MOD: u64, T: Into<GaloisField<MOD>>> $trait<T> for &GaloisField<MOD> {
                type Output = GaloisField<MOD>;
                fn $fn(self, rhs: T) -> Self::Output {
                    (*self).$fn(rhs.into())
                }
            }
            impl<const MOD: u64> $trait<&GaloisField<MOD>> for &GaloisField<MOD> {
                type Output = GaloisField<MOD>;
                fn $fn(self, rhs: &GaloisField<MOD>) -> Self::Output {
                    (*self).$fn(*rhs)
                }
            }
        )*};
    }
    gf_forward_ops! {
        Add, AddAssign, add, add_assign,
        Sub, SubAssign, sub, sub_assign,
        Mul, MulAssign, mul, mul_assign,
        Div, DivAssign, div, div_assign,
    }
    impl<const MOD: u64> Neg for GaloisField<MOD> {
        type Output = Self;
        fn neg(mut self) -> Self::Output {
            if self.value > 0 {
                self.value = MOD - self.value;
            }
            self
        }
    }
    impl<const MOD: u64> Sum for GaloisField<MOD> {
        fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
            iter.fold(Self::new(0), |acc, x| acc + x)
        }
    }
    impl<'a, const MOD: u64> Sum<&'a Self> for GaloisField<MOD> {
        fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
            iter.copied().sum()
        }
    }
    impl<const MOD: u64> Product for GaloisField<MOD> {
        fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
            iter.fold(Self::new(1), |acc, x| acc * x)
        }
    }
    impl<'a, const MOD: u64> Product<&'a Self> for GaloisField<MOD> {
        fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
            iter.copied().product()
        }
    }
}
0