結果

問題 No.1614 Majority Painting on Tree
コンテスト
ユーザー rhoo
提出日時 2025-07-15 18:02:41
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 1,257 ms / 5,000 ms
コード長 7,243 bytes
コンパイル時間 13,303 ms
コンパイル使用メモリ 402,336 KB
実行使用メモリ 14,976 KB
最終ジャッジ日時 2025-07-15 18:03:10
合計ジャッジ時間 27,826 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 4
other AC * 45
権限があれば一括ダウンロードができます

ソースコード

diff # 

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



#[fastout]
fn main(){
    input!{
        n:usize,
        c:usize,
        es:[(Usize1,Usize1);n-1],
    }

    let mut g=vec![vec![];n];
    for &(u,v) in &es{
        g[u].push(v);
        g[v].push(u);
    }

    let size=(0..n).map(|i|g[i].len() as u64).collect::<Vec<_>>();
    let pre=Pre::<M>::new(2*n+500);

    let get=|c:M,n:u64|->M{
        let mut ans=c.pow(n);
        for k in n/2+1..n{
            ans-=c*pre.choose(n as usize,k as usize)*(c-1).pow(n-k);
        }
        ans
    };

    let solve=|c:usize|->M{
        let c=M::new(c);
        let mut ans=get(c,size[0]);
        let c_inv=c.inv();
        
        for i in 1..g.len(){
            ans*=get(c,size[i])*c_inv;
        }

        ans
    };

    let mut ans=M::new(0);
    for i in 0..c{
        // i色使ってない色を確定する
        let sign=if i%2==0{
            M::new(1)
        } else{
            -M::new(1)
        };

        ans+=solve(c-i)*pre.choose(c,i)*sign;
    }

    println!("{ans}");
}



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;
}



// Modを超える数はpanicする
// 注意：
//      Modは素数
//      multi_choose(n,k)を使うときは，fac(n+k-1)まで必要になる
struct Pre<M:ModIntBase>{
    fac:Vec<M>,
    finv:Vec<M>,
}
impl<M:ModIntBase> Pre<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 inv(&self,n:usize)->M{
        assert!(n!=0);
        self.finv[n]*self.fac[n-1]
    }

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