#1075435 (Rust) No.3117 Reversible Tile

• 問題: No.3117 Reversible Tile
• ユーザー: rhoo
• 提出日時: 2025-04-18 22:49:00
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 403 ms / 3,000 ms
• コード長: 10,541 bytes
• コンパイル時間: 24,950 ms
• コンパイル使用メモリ: 392,516 KB
• 実行使用メモリ: 7,844 KB
• 最終ジャッジ日時: 2025-04-18 22:49:32
• 合計ジャッジ時間: 20,430 ms
• ジャッジサーバーID （参考情報）: judge1 / judge2

ファイルパターン	結果
sample	AC * 3
other	AC * 24

ソースコード

#![allow(unused_imports,non_snake_case,dead_code)]
use std::{cmp::Reverse as Rev,collections::*,iter::*};
use proconio::{marker::*,*};


#[fastout]
fn main(){
    input!{
        n:usize,
        m:usize,
        a:[usize;n],
    }

    let mut ones=a[0];
    for w in a.windows(2){
        ones+=w[0]^w[1];
    }

    let M=|a|M::new(a);
    let mut dp=vec![M(0);n+1];
    dp[0]+=1;

    for _ in 0..m{
        let mut new=vec![M(0);n+1];
        for i in 0..=n{
            if i!=0{
                new[i-1]+=dp[i]*i;
            }
            if i!=n{
                new[i+1]+=dp[i]*(n-i);
            }
            new[i]+=dp[i]*i*(n-i);
            if 2<=i{
                new[i-2]+=dp[i]*(i*(i-1)/2);
            }
            if i+2<=n{
                new[i+2]+=dp[i]*((n-i)*(n-i-1)/2);
            }
        }
        // eprintln!("{:?}",new);
        dp=new;
    }

    let com=Combination::<M>::new(n);

    println!("{}",dp[ones]/com.choose(n,ones));
}







type M = StaticModInt<998244353>;
type ModInt1000000007 = StaticModInt<1000000007>;



#[derive(Clone, Copy, PartialEq, Eq, Default, Hash)]
struct StaticModInt<const MODULO: u32> {
    val: u32,
}

impl<const MODULO: u32> ModIntBase for StaticModInt<MODULO> {
    fn modulus() -> u32 {
        MODULO
    }
    unsafe fn raw(val: u32) -> Self {
        StaticModInt { val }
    }
    fn val(self) -> u32 {
        self.val
    }
    fn inv(self) -> Self {
        assert!(self.val != 0, "divisor by zero");
        let inv_val = mod_inv_by_ext_gcd(self.val, MODULO);
        unsafe { Self::raw(inv_val) }
    }
}



impl<const MODULO: u32> std::fmt::Display for StaticModInt<MODULO> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.val)
    }
}

impl<const MODULO: u32> std::fmt::Debug for StaticModInt<MODULO> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.val)
    }
}


macro_rules! impl_from_int {
    ($($t:ty),*) => {
        $(impl<const MODULO: u32> From<$t> for StaticModInt<MODULO> {
            fn from(x: $t) -> Self {
                StaticModInt::new(x)
            }
        })*
    }
}

impl_from_int!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);



impl<const MODULO: u32> std::ops::Add for StaticModInt<MODULO> {
    type Output = Self;
    fn add(self, rhs: Self) -> Self::Output {
        let mut sum = self.val + rhs.val;
        if sum >= MODULO {
            sum -= MODULO;
        }
        unsafe { Self::raw(sum) }
    }
}

impl<const MODULO: u32> std::ops::Sub for StaticModInt<MODULO> {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self::Output {
        let diff = if self.val >= rhs.val {
            self.val - rhs.val
        } else {
            MODULO + self.val - rhs.val
        };
        unsafe { Self::raw(diff) }
    }
}

impl<const MODULO: u32> std::ops::Mul for StaticModInt<MODULO> {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self::Output {
        let prod = (self.val as u64 * rhs.val as u64 % MODULO as u64) as u32;
        unsafe { Self::raw(prod) }
    }
}

impl<const MODULO: u32> std::ops::Div for StaticModInt<MODULO> {
    type Output = Self;
    fn div(self, rhs: Self) -> Self::Output {
        self * rhs.inv()
    }
}



impl<const MODULO: u32> std::ops::AddAssign for StaticModInt<MODULO> {
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl<const MODULO: u32> std::ops::SubAssign for StaticModInt<MODULO> {
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl<const MODULO: u32> std::ops::MulAssign for StaticModInt<MODULO> {
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs;
    }
}

impl<const MODULO: u32> std::ops::DivAssign for StaticModInt<MODULO> {
    fn div_assign(&mut self, rhs: Self) {
        *self = *self / rhs;
    }
}



impl<const MODULO: u32> std::ops::Neg for StaticModInt<MODULO> {
    type Output = Self;
    fn neg(self) -> Self::Output {
        if self.val == 0 {
            self
        } else {
            unsafe { Self::raw(MODULO - self.val) }
        }
    }
}



impl<const MODULO: u32> std::str::FromStr for StaticModInt<MODULO> {
    type Err = <u32 as std::str::FromStr>::Err;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let x = s.parse::<u64>()?;
        Ok(StaticModInt::new(x))
    }
}



macro_rules! impl_mod_int_ops_sub {
    ($trait:ident, $func:ident, $assign_trait:ident, $assign_func:ident, $op:tt, ($($t:ty),*)) => {
        $(
            impl<const MODULO: u32> std::ops::$trait<$t> for StaticModInt<MODULO> {
                type Output = Self;
                fn $func(self, rhs: $t) -> Self::Output {
                    self $op StaticModInt::new(rhs)
                }
            }
            impl<const MODULO: u32> std::ops::$trait<StaticModInt<MODULO>> for $t {
                type Output = StaticModInt<MODULO>;
                fn $func(self, rhs: StaticModInt<MODULO>) -> Self::Output {
                    StaticModInt::new(self) $op rhs
                }
            }
            impl<const MODULO: u32> std::ops::$assign_trait<$t> for StaticModInt<MODULO> {
                fn $assign_func(&mut self, rhs: $t) {
                    *self = *self $op StaticModInt::new(rhs);
                }
            }
        )*
    }
}

macro_rules! impl_mod_int_ops {
    ($($t:tt)*) => {
        impl_mod_int_ops_sub!($($t)*, (i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize));
    }
}

impl_mod_int_ops!(Add, add, AddAssign, add_assign, +);
impl_mod_int_ops!(Sub, sub, SubAssign, sub_assign, -);
impl_mod_int_ops!(Mul, mul, MulAssign, mul_assign, *);
impl_mod_int_ops!(Div, div, DivAssign, div_assign, /);



impl<const MODULO: u32> std::iter::Sum for StaticModInt<MODULO> {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(StaticModInt::zero(), |acc, x| acc + x)
    }
}

impl<const MODULO: u32> std::iter::Product for StaticModInt<MODULO> {
    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(StaticModInt::one(), |acc, x| acc * x)
    }
}



trait RemEuclidU32 {
    fn rem_euclid_u32(self, modulus: u32) -> u32;
}

macro_rules! impl_rem_euclid_u32_for_small_signed {
    ($($ty:tt),*) => {
        $(
            impl RemEuclidU32 for $ty {
                fn rem_euclid_u32(self, modulus: u32) -> u32 {
                    (self as i64).rem_euclid(i64::from(modulus)) as _
                }
            }
        )*
    }
}

impl_rem_euclid_u32_for_small_signed!(i8, i16, i32, i64, isize);

impl RemEuclidU32 for i128 {
    fn rem_euclid_u32(self, modulus: u32) -> u32 {
        self.rem_euclid(i128::from(modulus)) as _
    }
}

macro_rules! impl_rem_euclid_u32_for_small_unsigned {
    ($($ty:tt),*) => {
        $(
            impl RemEuclidU32 for $ty {
                fn rem_euclid_u32(self, modulus: u32) -> u32 {
                    self as u32 % modulus
                }
            }
        )*
    }
}

macro_rules! impl_rem_euclid_u32_for_large_unsigned {
    ($($ty:tt),*) => {
        $(
            impl RemEuclidU32 for $ty {
                fn rem_euclid_u32(self, modulus: u32) -> u32 {
                    (self % (modulus as $ty)) as _
                }
            }
        )*
    }
}

impl_rem_euclid_u32_for_small_unsigned!(u8, u16, u32);
impl_rem_euclid_u32_for_large_unsigned!(u64, u128);


#[cfg(target_pointer_width = "32")]
impl_rem_euclid_u32_for_small_unsigned!(usize);

#[cfg(target_pointer_width = "64")]
impl_rem_euclid_u32_for_large_unsigned!(usize);



trait ModIntBase: Default
    + std::str::FromStr
    + From<i8>
    + From<i16>
    + From<i32>
    + From<i64>
    + From<i128>
    + From<isize>
    + From<u8>
    + From<u16>
    + From<u32>
    + From<u64>
    + From<u128>
    + From<usize>
    + Copy
    + Eq
    + std::hash::Hash
    + std::fmt::Display
    + std::fmt::Debug
    + std::ops::Neg<Output = Self>
    + std::ops::Add<Output = Self>
    + std::ops::Sub<Output = Self>
    + std::ops::Mul<Output = Self>
    + std::ops::Div<Output = Self>
    + std::ops::AddAssign
    + std::ops::SubAssign
    + std::ops::MulAssign
    + std::ops::DivAssign
{
    fn modulus() -> u32;
    unsafe fn raw(val: u32) -> Self;
    fn val(self) -> u32;
    fn inv(self) -> Self;

    fn new(val: impl RemEuclidU32) -> Self {
        unsafe { Self::raw(val.rem_euclid_u32(Self::modulus())) }
    }

    fn zero() -> Self {
        unsafe { Self::raw(0) }
    }

    fn one() -> Self {
        unsafe { Self::raw(1) }
    }

    fn pow(self, mut n: u64) -> Self {
        let mut x = self;
        let mut r = Self::one();
        while n > 0 {
            if n & 1 == 1 {
                r *= x;
            }
            x *= x;
            n >>= 1;
        }
        r
    }
}



fn mod_inv_by_ext_gcd(x: u32, modulus: u32) -> u32 {
    let (mut a, mut b) = (x as i64, modulus as i64);
    let (mut u, mut v) = (1i64, 0i64);
    while b != 0 {
        let t = a / b;

        a -= t * b;
        (a,b)=(b,a);

        u -= t * v;
        (u,v)=(v,u);
    }
    assert!(a==1);

    ((u % modulus as i64 + modulus as i64) % modulus as i64) as u32
}



// Modを超える数はpanicする
// 注意：
//      Modは素数
//      multi_choose(n,k)を使うときは，fac(n+k-1)まで必要になる
struct Combination<M:ModIntBase>{
    fac:Vec<M>,
    finv:Vec<M>,
}
impl<M:ModIntBase> Combination<M>{
    fn new(n:usize)->Self{
        assert!(n<M::modulus() as usize);

        let mut fac=vec![M::new(1);n+1];
        for i in 1..=n{
            fac[i]=fac[i-1]*M::new(i);
        }

        let mut finv=vec![M::new(0);n+1];
        finv[n]=fac[n].inv();
        for i in (1..=n).rev(){
            finv[i-1]=finv[i]*M::new(i);
        }
        
        Self{fac,finv}
    }

    fn fac(&self,n:usize)->M{
        self.fac[n]
    }

    fn finv(&self,n:usize)->M{
        self.finv[n]
    }

    fn perm(&self,n:usize,k:usize)->M{
        if n<k || (n as i64)<0 || (k as i64)<0{
            return M::new(0);
        }
        self.fac(n)*self.finv(n-k)
    }

    fn choose(&self,n:usize,k:usize)->M{
        if n<k || (n as i64)<0 || (k as i64)<0{
            return M::new(0);
        }
        self.fac(n)*self.finv(k)*self.finv(n-k)
    }

    fn multi_choose(&self,n:usize,k:usize)->M{
        if (n as i64)<0 || (k as i64)<0{
            return M::new(0);
        }
        if n==0 && k==0{
            return M::new(1);
        }
        self.choose(n+k-1,k)
    }
}