#1131015 (Rust) No.3370 AB → BA

提出ソース

結果

問題	No.3370 AB → BA
コンテスト
ユーザー	rhoo
提出日時	2025-11-03 00:48:15
言語	Rust (1.83.0 + proconio)
結果	WA
実行時間	-
コード長	12,855 bytes
コンパイル時間	13,487 ms
コンパイル使用メモリ	400,424 KB
実行使用メモリ	26,072 KB
最終ジャッジ日時	2025-11-17 20:43:13
合計ジャッジ時間	24,750 ms
ジャッジサーバーID （参考情報）	judge5 / judge2

このコードへのチャレンジ
（要ログイン）

ファイルパターン	結果
sample	AC * 1
other	AC * 10 WA * 10

権限があれば一括ダウンロードができます

コンパイルメッセージ

warning: type alias `ModInt1000000007` is never used
  --> src/main.rs:85:6
   |
85 | type ModInt1000000007=ModInt<1000000007>;
   |      ^^^^^^^^^^^^^^^^
   |
   = note: `#[warn(dead_code)]` on by default

warning: method `val` is never used
   --> src/main.rs:268:8
    |
262 | trait ModIntBase:Default+std::str::FromStr+Copy+Eq+std::hash::Hash+std::fmt::Display+std::fmt::Debug
    |       ---------- method in this trait
...
268 |     fn val(self)->u32;
    |        ^^^

ソースコード

raw source code

#![allow(non_snake_case)]



fn main(){
    proconio::input!{
        s:proconio::marker::Chars,
    }

    let mut a=vec![];
    let mut cnt=1;
    for &c in &s{
        if c=='B'{
            cnt+=1;
        } else{
            a.push(cnt);
        }
    }
    a.push(cnt);

    if a.is_empty(){
        println!("1");
        return;
    }

    let n=s.len()*2+1;
    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);
    }
    
    cap!{
        #![&fac:[M],&finv:[M]]
        fn rec(a:&[usize],D:&[M],add:usize)->Vec<M>{
            if a.len()==1{
                return vec![D[0];a[0]-add];
            }
            
            let mi=a.len()/2;
            if a[mi]==add{
                return rec!(&a[mi..],&D[mi..],add);
            }

            let mut L=rec!(&a[..mi],&D[..mi],add);
            L.resize(a[mi]-add,M::new(0));
            let D=&D[mi..];

            let (n,m)=(L.len()-1,D.len()-1);
            let mut R=L.conv(&(0..=n).map(|i|fac[m+i]*finv[i]*finv[m]).collect::<Vec<_>>());
            let mut U=D.conv(&(0..=m).map(|i|fac[n+i]*finv[i]*finv[m]).collect::<Vec<_>>());
            let addR=fac[..=n+m].conv(&(0..=m).map(|i|D[i]*finv[m-i]).collect::<Vec<_>>());
            let addU=fac[..=n+m].conv(&(0..=n).map(|i|L[i]*finv[n-i]).collect::<Vec<_>>());
            R.truncate(n+1);
            U.truncate(m+1);

            for i in 0..=n{
                R[i]+=addR[m+i]*finv[i];
            }
            for i in 0..=m{
                U[i]+=addU[n+i]*finv[i];
            }
            R.extend(&rec!(&a[mi..],&U,a[mi]));
            R
        }
    }

    let mut D=vec![M::new(0);a.len()];
    D[0]+=1;
    let res=rec!(&a,&D,0);
    println!("{}",res.last().unwrap());
}



type M=ModInt998244353;



type ModInt998244353=ModInt<998244353>;
type ModInt1000000007=ModInt<1000000007>;



#[derive(Clone,Copy,PartialEq,Eq,Default,Hash)]
struct ModInt<const MOD:u32>(u32);
impl<const MOD:u32> ModIntBase for ModInt<MOD>{
    fn modulus()->u32{ MOD }
    fn val(self)->u32{ self.0 }
    fn new(v:impl RemU32)->Self{ Self(v.rem_u32(MOD)) }

    fn inv(self)->Self{
        assert!(self.0!=0);

        let (mut a,mut b)=(self.0 as i64,MOD as i64);
        let (mut u,mut v)=(1,0);
        while b!=0{
            let t=a/b;
            (a,b)=(b,a-t*b);
            (u,v)=(v,u-t*v);
        }
        assert!(a==1);

        if u<0{
            u+=MOD as i64;
        }
        Self(u as u32)
    }

    fn pow(self,mut k:u64)->Self{
        let mut pow2=self;
        let mut ret=Self(1);
        while k>0{
            if k&1==1{
                ret*=pow2;
            }
            pow2*=pow2;
            k>>=1;
        }
        ret
    }
}


impl<const MOD:u32> std::fmt::Display for ModInt<MOD>{
    fn fmt(&self,f:&mut std::fmt::Formatter)->std::fmt::Result{
        write!(f,"{}",self.0)
    }
}
impl<const MOD:u32> std::fmt::Debug for ModInt<MOD>{
    fn fmt(&self,f:&mut std::fmt::Formatter)->std::fmt::Result{
        write!(f,"{}",self.0)
    }
}


impl<const MOD:u32> std::ops::Add for ModInt<MOD>{
    type Output=Self;
    fn add(self,a:Self)->Self{
        let mut new=self.0+a.0;
        if MOD<=new{
            new-=MOD;
        }
        Self(new)
    }
}
impl<const MOD:u32> std::ops::Sub for ModInt<MOD>{
    type Output=Self;
    fn sub(self,a:Self)->Self{
        let mut new=self.0-a.0;
        if 0>new as i32{
            new+=MOD;
        }
        Self(new)
    }
}
impl<const MOD:u32> std::ops::Mul for ModInt<MOD>{
    type Output=Self;
    fn mul(self,a:Self)->Self{
        Self((self.0 as u64*a.0 as u64%MOD as u64) as u32)
    }
}
impl<const MOD:u32> std::ops::Div for ModInt<MOD>{
    type Output=Self;
    fn div(self,a:Self)->Self{
        self*a.inv()
    }
}
impl<const MOD:u32> std::ops::Neg for ModInt<MOD>{
    type Output=Self;
    fn neg(self)->Self{
        if self.0==0{
            return self;
        }
        Self(MOD-self.0)
    }
}


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


macro_rules! impl_modint_ops{
    ($trait:ident,$func:ident,$assign_trait:ident,$assign_func:ident,$op:tt)=>{
        impl<const MOD:u32> std::ops::$assign_trait for ModInt<MOD>{
            fn $assign_func(&mut self,a:Self){
                *self=*self $op a
            }
        }
        impl<T:RemU32,const MOD:u32> std::ops::$trait<T> for ModInt<MOD>{
            type Output=Self;
            fn $func(self,a:T)->Self{
                self $op Self::new(a)
            }
        }
        impl<T:RemU32,const MOD:u32> std::ops::$assign_trait<T> for ModInt<MOD>{
            fn $assign_func(&mut self,a:T){
                *self=*self $op Self::new(a)
            }
        }
    }
}
impl_modint_ops!(Add,add,AddAssign,add_assign,+);
impl_modint_ops!(Sub,sub,SubAssign,sub_assign,-);
impl_modint_ops!(Mul,mul,MulAssign,mul_assign,*);
impl_modint_ops!(Div,div,DivAssign,div_assign,/);


impl<const MOD:u32> std::iter::Sum for ModInt<MOD>{
    fn sum<I:Iterator<Item=Self>>(iter:I)->Self{
        iter.fold(Self(0),|sum,x|sum+x)
    }
}
impl<const MOD:u32> std::iter::Product for ModInt<MOD>{
    fn product<I:Iterator<Item=Self>>(iter:I)->Self{
        iter.fold(Self(1),|prod,x|prod*x)
    }
}


trait RemU32{
    fn rem_u32(self,m:u32)->u32;
}
macro_rules! impl_rem_u32{
    ($($ty:tt),*)=>{
        $(
            impl RemU32 for $ty{
                fn rem_u32(self,m:u32)->u32{
                    (self as i64).rem_euclid(m as i64) as _
                }
            }
        )*
    }
}
impl_rem_u32!(i32,i64);


macro_rules! impl_rem_u32{
    ($($ty:tt),*)=>{
        $(
            impl RemU32 for $ty{
                fn rem_u32(self,m:u32)->u32{
                    (self%(m as $ty)) as _
                }
            }
        )*
    }
}
impl_rem_u32!(u32,u64,usize);


trait ModIntBase:Default+std::str::FromStr+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;
    fn val(self)->u32;
    fn new(v:impl RemU32)->Self;
    fn inv(self)->Self;
    fn pow(self,k:u64)->Self;
}



type NttMod=ModInt998244353;


trait Ntt{
    fn ntt(&mut self);
    fn intt(&mut self);
    fn conv(&self,a:&[NttMod])->Vec<NttMod>;
}


impl Ntt for [NttMod]{
    // bitreverse順になる
    fn ntt(&mut self){
        let n=self.len();
        assert!(NttMod::modulus()==998244353);
        assert!(n.is_power_of_two());
        assert!(n<=1<<(NttMod::modulus()-1).trailing_zeros());

        let len=n.trailing_zeros() as usize;
        let sum_e=&get_ntt_pre().0;

        for ph in 1..=len{
            let p=1<<len-ph;
            let mut now=NttMod::new(1);

            for (i,f) in self.chunks_exact_mut(2*p).enumerate(){
                let (x,y)=f.split_at_mut(p);
                for (x,y) in std::iter::zip(x,y){
                    let r=*y*now;
                    (*x,*y)=(*x+r,*x-r);
                }
                now*=sum_e[i.trailing_ones() as usize];
            }
        }
    }

    // bitreverse順になる
    fn intt(&mut self){
        let n=self.len();
        assert!(NttMod::modulus()==998244353);
        assert!(n.is_power_of_two());
        assert!(n<=1<<(NttMod::modulus()-1).trailing_zeros());

        let len=n.trailing_zeros() as usize;
        let sum_ie=&get_ntt_pre().1;

        for ph in (1..=len).rev(){
            let p=1<<len-ph;
            let mut inow=NttMod::new(1);

            for (i,f) in self.chunks_exact_mut(2*p).enumerate(){
                let (x,y)=f.split_at_mut(p);
                for (x,y) in std::iter::zip(x,y){
                    (*x,*y)=(*x+*y,(*x-*y)*inow);
                }
                inow*=sum_ie[i.trailing_ones() as usize];
            }
        }

        let ik=NttMod::new(1<<len).inv();
        self.iter_mut().for_each(|f|*f*=ik);
    }

    fn conv(&self,a:&[NttMod])->Vec<NttMod>{
        let (n,m)=(self.len(),a.len());
        
        if n.min(m)<=60{
            if n==0 || m==0{
                return vec![];
            }

            let mut res=vec![NttMod::new(0);n+m-1];
            for (i,x) in self.iter().enumerate(){
                for (res,&y) in std::iter::zip(&mut res[i..],a){
                    *res+=*x*y;
                }
            }

            return res;
        }

        let size=(n+m-1).next_power_of_two();
        let mut f=self.to_vec();
        f.resize(size,NttMod::new(0));
        let mut g=a.to_vec();
        g.resize(size,NttMod::new(0));

        f.ntt();
        g.ntt();
        f.iter_mut().zip(&g).for_each(|(f,&g)|*f*=g);

        f.intt();
        f.truncate(n+m-1);
        f
    }
}



fn get_ntt_pre()->&'static ([NttMod;30],[NttMod;30]){
    use std::sync::OnceLock;
    static NTT_PRE:OnceLock<([NttMod;30],[NttMod;30])>=OnceLock::new();

    NTT_PRE.get_or_init(||{
        let cnt2=(NttMod::modulus()-1).trailing_zeros() as usize;

        let mut es=[NttMod::new(0);30];
        let (mut ies,mut sum_e,mut sum_ie)=(es,es,es);

        let mut e=NttMod::new(3).pow((NttMod::modulus()-1) as u64>>cnt2);
        let mut ie=e.inv();

        for i in (2..=cnt2).rev(){
            es[i-2]=e;
            ies[i-2]=ie;
            e*=e;
            ie*=ie;
        }

        let mut now=NttMod::new(1);
        let mut inow=NttMod::new(1);

        for i in 0..cnt2-1{
            sum_e[i]=es[i]*now;
            sum_ie[i]=ies[i]*inow;
            now*=ies[i];
            inow*=es[i];
        }

        (sum_e,sum_ie)
    })
}



#[macro_export]
macro_rules! cap{
    ()=>{};
    (#![] $($fn:tt)*)=>{
        cap!([],[],[],[],[$($fn)*],[],[]);
    };
    (#![$($g:tt)*] $($fn:tt)*)=>{
        cap!([$($g)*,],[],[],[],[$($fn)*],[],[]);
    };
    ([$name:ident:$t:ty,$($rem:tt)*],[$($ga:tt)*],[$($gn:tt)*],[$($ge:tt)*],[$($fn:tt)*],[],[])=>{
        cap!([$($rem)*],[$name:$t,$($ga)*],[$name,$($gn)*],[$name,$($ge)*],[$($fn)*],[],[]);
    };
    ([$($flag:tt)? $name:ident:$t:ty,$($rem:tt)*],[$($ga:tt)*],[$($gn:tt)*],[$($ge:tt)*],[$($fn:tt)*],[],[])=>{
        cap!([$($rem)*],[$name:$($flag)?$t,$($ga)*],[$name,$($gn)*],[$($flag)?$name,$($ge)*],[$($fn)*],[],[]);
    };
    ([&mut $name:ident:$t:ty,$($rem:tt)*],[$($ga:tt)*],[$($gn:tt)*],[$($ge:tt)*],[$($fn:tt)*],[],[])=>{
        cap!([$($rem)*],[$name:&mut $t,$($ga)*],[$name,$($gn)*],[&mut $name,$($ge)*],[$($fn)*],[],[]);
    };
    ([$(,)?],[$($ga:tt)*],[$($gn:tt)*],[$($ge:tt)*],[$(#[$($att:tt)*])? fn $name:ident($($arg:tt)*) $body:block $($rem:tt)*],[$($info:tt)*],[$($fn:tt)*])=>{
        cap!([],[$($ga)*],[$($gn)*],[$($ge)*],[$($rem)*],[(($(#[$($att)*])?),$name,($($arg)*),(),$body),$($info)*],[$name,$($fn)*]);
    };
    ([$(,)?],[$($ga:tt)*],[$($gn:tt)*],[$($ge:tt)*],[$(#[$($att:tt)*])? fn $name:ident($($arg:tt)*)->$ret:ty $body:block $($rem:tt)*],[$($info:tt)*],[$($fn:tt)*])=>{
        cap!([],[$($ga)*],[$($gn)*],[$($ge)*],[$($rem)*],[(($(#[$($att)*])?),$name,($($arg)*),$ret,$body),$($info)*],[$name,$($fn)*]);
    };
    ([$(,)?],[$($ga:tt)*],[$($gn:tt)*],[$($ge:tt)*],[],[$($info:tt)*],[$($fn:tt)*])=>{
        cap!(@make_fn [$($ga)*],[$($gn)*],[$($ge)*],[$($info)*],[$($fn)*]);
    };
    (@make_fn [$($ga:ident:$gt:ty,)*],[$($gn:tt)*],[$($ge:tt)*],[(($($att:tt)*),$name:ident,($($arg:tt)*),$ret:ty,$body:block),$($rem:tt)*],[$($fn:tt)*])=>{
        $($att)*
        fn $name($($ga:$gt,)*$($arg)*)->$ret{
            $(#[allow(unused_variables)] let $ga=$ga;)*
            cap!(@make_macros ($),[$($gn)*],[$($fn)*]);
            $body
        }
        cap!(@make_fn [$($ga:$gt,)*],[$($gn)*],[$($ge)*],[$($rem)*],[$($fn)*]);
    };
    (@make_fn [$($ga:tt)*],[$($gn:tt)*],[$($ge:tt)*],[],[$($fn:tt)*])=>{
        cap!(@make_global_macros ($),[$($ge)*],[$($fn)*]);
    };
    (@make_macros ($dol:tt),[$($gn:ident,)*],[$name:ident,$($rem:tt)*])=>{
        #[allow(unused_macros)]
        macro_rules! $name{
            ($dol($dol arg:expr),*)=>{$name($($gn,)* $dol($dol arg),*)}
        }
        cap!(@make_macros ($),[$($gn,)*],[$($rem)*]);
    };
    (@make_macros ($dol:tt),[$($gn:ident,)*],[])=>{};
    (@make_global_macros ($dol:tt),[$($ge:expr,)*],[$name:ident,$($rem:tt)*])=>{
        #[allow(unused_macros)]
        macro_rules! $name{
            ($dol($dol arg:expr),*)=>{$name($($ge,)* $dol($dol arg),*)}
        }
        cap!(@make_global_macros ($),[$($ge,)*],[$($rem)*]);
    };
    (@make_global_macros ($dol:tt),[$($ge:expr,)*],[])=>{};
}

yukicoder

結果

コンパイルメッセージ

ソースコード