結果

問題 No.3343 Distance Sum of Large Tree
コンテスト
ユーザー rhoo
提出日時 2025-11-13 22:18:01
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 64 ms / 2,000 ms
コード長 6,328 bytes
コンパイル時間 14,090 ms
コンパイル使用メモリ 395,224 KB
実行使用メモリ 24,192 KB
最終ジャッジ日時 2025-11-13 22:18:18
合計ジャッジ時間 16,096 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 30
権限があれば一括ダウンロードができます

ソースコード

diff # 

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



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

    let mut g=vec![vec![];n];
    let mut pot=vec![vec![];n];

    for i in 1..n{
        let u=p[i-1];
        let v=i;
        
        g[u].push(v);
        pot[u].push((c[i-1],v));
        pot[v].push((b[i-1],u));
    }
    let mut sub=vec![0;n];
    for v in (0..n).rev(){
        sub[v]=a[v];
        for &nv in &g[v]{
            sub[v]+=sub[nv];
        }
    }

    let tot=sub[0];
    let get=|u:usize,v:usize|->usize{
        assert!(u!=v);
        if u<v{
            sub[v]
        } else{
            tot-sub[u]
        }
    };

    let mut ans=M::new(0);
    for v in 0..n{
        for &nv in &g[v]{
            ans+=M::new(1)*get(v,nv)*get(nv,v);
        }
    }

    for u in 0..n{
        pot[u].sort();
        let mut rem=1;
        let mut i=0;

        for &(ni,v) in &pot[u]{
            if i==ni{
                rem+=get(u,v);
            } else{
                let L=rem;
                let M=ni-i-1;
                let R=tot-L-M;

                ans+=M::new(1)*(M+1)*L*R;
                ans+=M::new(1)*(L+R)*(M+1)*M/2;
                ans+=M::new(1)*(M+1)*M*(M-1)/6;

                rem+=ni-i+get(u,v);
                i=ni;
            }
        }

        let L=rem;
        let M=tot-L;
        ans+=M::new(1)*L*(M+1)*M/2;
        ans+=M::new(1)*(M+1)*M*(M-1)/6;
    }

    println!("{}",ans*2);
}


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