結果

問題 No.2618 除霊
ユーザー 37kt
提出日時 2024-12-02 15:10:09
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 512 ms / 2,000 ms
コード長 10,611 bytes
コンパイル時間 12,132 ms
コンパイル使用メモリ 379,004 KB
実行使用メモリ 68,860 KB
最終ジャッジ日時 2024-12-02 15:10:48
合計ジャッジ時間 33,217 ms
ジャッジサーバーID
(参考情報)
judge5 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
other AC * 43
権限があれば一括ダウンロードができます

ソースコード

diff #
プレゼンテーションモードにする

pub use __cargo_equip::prelude::*;
use algebraic::{act, algebra, monoid};
use graph::Graph;
#[allow(unused_imports)]
use proconio::{
input,
marker::{Bytes, Chars, Usize1},
};
use re_rooting_dp::ReRootingDP;
#[derive(Clone, Copy, Default)]
struct S {
f: [usize; 3],
s: [usize; 2],
g: [bool; 2],
}
algebra!(M, S);
monoid!(M, S::default(), |&a: &S, &b: &S| {
S {
f: [a.f[0] + b.f[0], a.f[1] + b.f[1], a.f[2] + b.f[2]],
s: [a.s[0] + b.s[0], a.s[1] + b.s[1]],
g: [a.g[0] || b.g[0], a.g[1] || b.g[1]],
}
});
algebra!(V, bool);
act!(V, S, |&v, &a: &S| {
let mut a = a;
a.s[0] = 1;
a.g[0] = v;
if a.g[1] {
a.f[1] += a.s[0];
a.s[0] = 0;
}
if a.g[0] {
a.f[0] += a.s[0] + a.s[1];
a.s[0] = 0;
a.s[1] = 0;
}
a
});
algebra!(E, ());
act!(E, S, |_, &a: &S| {
S {
f: [0, a.f[0], a.f[1] + a.f[2]],
s: [0, a.s[0]],
g: [false, a.g[0]],
}
});
fn main() {
input! {
n: usize,
ab: [(Usize1, Usize1); n - 1],
v: [Usize1],
}
let mut f = vec![false; n];
for &v in &v {
f[v] = true;
}
let g = Graph::from_vertices_and_unweighted_undirected_edges(&f, &ab);
let dp = ReRootingDP::build::<M, V, E>(&g);
for i in 0..n {
let res = dp.prod(i).f[2];
println!("{}", res);
}
}
// The following code was expanded by `cargo-equip`.
/// # Bundled libraries
///
/// - `algebraic 0.1.0 (path+███████████████████████████████████████████████)` published in **missing** licensed under `CC0-1.0` as `crate
    ::__cargo_equip::crates::algebraic`
/// - `graph 0.1.0 (path+███████████████████████████████████████)` published in **missing** licensed under `CC0-1.0` as `crate
    ::__cargo_equip::crates::graph`
/// - `re-rooting-dp 0.1.0 (path+██████████████████████████████████████████████)` published in **missing** licensed under `CC0-1.0` as `crate
    ::__cargo_equip::crates::re_rooting_dp`
#[cfg_attr(any(), rustfmt::skip)]
#[allow(unused)]
mod __cargo_equip {
pub(crate) mod crates {
pub mod algebraic {pub use crate::__cargo_equip::macros::algebraic::*;pub trait Algebra{type S;}pub trait Act:Algebra{type X;fn act(f:&Self
            ::S,x:&Self::X)->Self::X;}pub trait Monoid:Algebra{fn e()->Self::S;fn op(x:&Self::S,y:&Self::S)->Self::S;}pub trait Group:Monoid{fn inv(x
            :&Self::S)->Self::S;}pub trait Zero{fn zero()->Self;fn is_zero(&self)->bool;}pub trait One{fn one()->Self;fn is_one(&self)->bool
            ;}#[macro_export]macro_rules!__cargo_equip_macro_def_algebraic_algebra{($ident:ident,$ty:ty)=>{#[derive(Clone)]enum$ident{}impl$crate
            ::__cargo_equip::crates::algebraic::Algebra for$ident{type S=$ty;}};}macro_rules!algebra{($($tt:tt)*)=>(crate
            ::__cargo_equip_macro_def_algebraic_algebra!{$($tt)*})}#[macro_export]macro_rules!__cargo_equip_macro_def_algebraic_act{($ident:ident
            ,$tar:ty,$act:expr)=>{impl$crate::__cargo_equip::crates::algebraic::Act for$ident{type X=$tar;#[inline]fn act(f:&Self::S,x:&Self::X
            )->Self::X{$act(f,x)}}};}macro_rules!act{($($tt:tt)*)=>(crate::__cargo_equip_macro_def_algebraic_act!{$($tt)*}
            )}#[macro_export]macro_rules!__cargo_equip_macro_def_algebraic_monoid{($ident:ident,$e:expr,$op:expr)=>{impl$crate::__cargo_equip::crates
            ::algebraic::Monoid for$ident{#[inline]fn e()->Self::S{$e}#[inline]fn op(x:&Self::S,y:&Self::S)->Self::S{$op(x,y)}}};}macro_rules!monoid{
            ($($tt:tt)*)=>(crate::__cargo_equip_macro_def_algebraic_monoid!{$($tt)*}
            )}#[macro_export]macro_rules!__cargo_equip_macro_def_algebraic_group{($ident:ident,$e:expr,$op:expr,$inv:expr)=>{impl$crate
            ::__cargo_equip::crates::algebraic::Monoid for$ident{#[inline]fn e()->Self::S{$e}#[inline]fn op(x:&Self::S,y:&Self::S)->Self::S{$op(x,y
            )}}impl$crate::__cargo_equip::crates::algebraic::Group for$ident{#[inline]fn inv(x:&Self::S)->Self::S{$inv(x)}}};}macro_rules!group{($
            ($tt:tt)*)=>(crate::__cargo_equip_macro_def_algebraic_group!{$($tt)*})}macro_rules!impl_zero_one{($($t:ty)*)=>{$(impl$crate
            ::__cargo_equip::crates::algebraic::Zero for$t{fn zero()->Self{0}fn is_zero(&self)->bool{*self==0}}impl$crate::__cargo_equip::crates
            ::algebraic::One for$t{fn one()->Self{1}fn is_one(&self)->bool{*self==1}})*};}impl_zero_one!(usize u8 u16 u32 u64 u128 isize i8 i16 i32
            i64 i128);}
pub mod graph {use std::ops::Index;#[derive(Clone)]pub struct Graph<V,E>where V:Clone,E:Clone,{vertices:Vec<V>,edges:Vec<(usize,E)>,pos:Vec
            <usize>,edge_id:Vec<usize>,edges_count:usize,}pub const GRID_NEIGHBOURS_4:&[(usize,usize)]=&[(!0,0),(0,!0),(1,0),(0,1)];pub const
            GRID_NEIGHBOURS_8:&[(usize,usize)]=&[(!0,0),(0,!0),(1,0),(0,1),(!0,!0),(!0,1),(1,!0),(1,1),];impl<V,E>Graph<V,E>where V:Clone,E:Clone
            ,{pub fn from_vertices_and_directed_edges(vertices:&[V],edges:&[(usize,usize,E)])->Self{if edges.is_empty(){return Self{vertices:vertices
            .to_vec(),edges:vec![],pos:vec![0;vertices.len()+1],edge_id:vec![],edges_count:0,};}let n=vertices.len();let mut pos=vec![0;n+1];for&(u,_
            ,_)in edges{pos[u]+=1;}for i in 1..=n{pos[i]+=pos[i-1];}let mut ord=vec![0;edges.len()];for i in(0..edges.len()).rev(){let u=edges[i].0
            ;pos[u]-=1;ord[pos[u]]=i;}Self{vertices:vertices.to_vec(),edge_id:ord.clone(),edges:ord.into_iter().map(|i|(edges[i].1,edges[i].2.clone
            ())).collect(),pos,edges_count:edges.len(),}}pub fn from_vertices_and_undirected_edges(vertices:&[V],edges:&[(usize,usize,E)])->Self{if
            edges.is_empty(){return Self{vertices:vertices.to_vec(),edges:vec![],pos:vec![0;vertices.len()+1],edge_id:vec![],edges_count:0,};}let n
            =vertices.len();let mut pos=vec![0;n+1];for&(u,v,_)in edges{pos[u]+=1;pos[v]+=1;}for i in 1..=n{pos[i]+=pos[i-1];}let mut ord=vec![0
            ;edges.len()*2];for i in(0..edges.len()*2).rev(){if i&1==0{let u=edges[i>>1].0;pos[u]-=1;ord[pos[u]]=i;}else{let v=edges[i>>1].1;pos[v]
            -=1;ord[pos[v]]=i;}}Self{vertices:vertices.to_vec(),edge_id:ord.iter().map(|&i|i/2).collect(),edges:ord.into_iter().map(|i|{(if i&1
            ==0{edges[i>>1].1}else{edges[i>>1].0},edges[i>>1].2.clone(),)}).collect(),pos,edges_count:edges.len(),}}pub fn len(&self)->usize{self
            .vertices.len()}pub fn edges_count(&self)->usize{self.edges_count}pub fn vertex(&self,v:usize)->&V{&self.vertices[v]}pub fn edges(&self,v
            :usize)->&[(usize,E)]{let l=self.pos[v];let r=self.pos[v+1];&self.edges[l..r]}pub fn edge_id(&self,v:usize,i:usize)->usize{self
            .edge_id[self.pos[v]+i]}pub fn from_grid(grid:&Vec<Vec<V>>,neighbours:&[(usize,usize)],cost:impl Fn(&V,&V)->Option<E>,)->Self{let h=grid
            .len();let w=grid[0].len();let mut edges=vec![];for i in 0..h{for j in 0..w{for&(di,dj)in neighbours{let ni=i.wrapping_add(di);let nj=j
            .wrapping_add(dj);if ni>=h||nj>=w{continue;}if let Some(c)=cost(&grid[i][j],&grid[ni][nj]){edges.push((i*w+j,ni*w+nj,c));}}}}Self
            ::from_vertices_and_directed_edges(&grid.into_iter().flatten().cloned().collect::<Vec<_>>(),&edges,)}}impl<V,E>Index<usize>for Graph<V,E
            >where V:Clone,E:Clone,{type Output=[(usize,E)];fn index(&self,v:usize)->&[(usize,E)]{self.edges(v)}}impl<E>Graph<(),E>where E:Clone,{pub
            fn from_directed_edges(n:usize,edges:&[(usize,usize,E)])->Self{Self::from_vertices_and_directed_edges(&vec![();n],edges)}pub fn
            from_undirected_edges(n:usize,edges:&[(usize,usize,E)])->Self{Self::from_vertices_and_undirected_edges(&vec![();n],edges)}}impl<V>Graph<V
            ,()>where V:Clone,{pub fn from_vertices_and_unweighted_directed_edges(vertices:&[V],edges:&[(usize,usize)],)->Self{Self
            ::from_vertices_and_directed_edges(vertices,&edges.iter().map(|&(u,v)|(u,v,())).collect::<Vec<_>>(),)}pub fn
            from_vertices_and_unweighted_undirected_edges(vertices:&[V],edges:&[(usize,usize)],)->Self{Self::from_vertices_and_undirected_edges
            (vertices,&edges.iter().map(|&(u,v)|(u,v,())).collect::<Vec<_>>(),)}}impl Graph<(),()>{pub fn from_unweighted_directed_edges(n:usize
            ,edges:&[(usize,usize)])->Self{Self::from_directed_edges(n,&edges.iter().map(|&(u,v)|(u,v,())).collect::<Vec<_>>(),)}pub fn
            from_unweighted_undirected_edges(n:usize,edges:&[(usize,usize)])->Self{Self::from_undirected_edges(n,&edges.iter().map(|&(u,v)|(u,v
            ,())).collect::<Vec<_>>(),)}}}
pub mod re_rooting_dp {use crate::__cargo_equip::preludes::re_rooting_dp::*;use algebraic::{Act,Monoid};use graph::Graph;pub struct
            ReRootingDP<S>where S:Clone,{par:Vec<usize>,dp:Vec<S>,dpl:Vec<S>,dph:Vec<S>,}impl<T>ReRootingDP<T>where T:Clone,{pub fn build<M,V,E>(g
            :&Graph<V::S,E::S>)->Self where M:Monoid<S=T>,M::S:Clone,V:Act<X=M::S>,V::S:Clone,E:Act<X=M::S>,E::S:Clone,{let n=g.len();let mut ord
            =vec![];let mut par=vec![!0;n];let mut st=vec![0];while let Some(v)=st.pop(){ord.push(v);for&(u,_)in&g[v]{if u!=0&&par[u]==!0{par[u]=v;st
            .push(u);}}}let mut dpl=vec![M::e();n];let mut dph=vec![M::e();n];let mut dp=vec![M::e();n];for&v in ord.iter().rev(){let m=g[v].len
            ();let mut xl=vec![M::e();m+1];let mut xr=vec![M::e();m+1];for i in 0..m{let(u,_)=g[v][i];if u==par[v]{xl[i+1]=xl[i].clone();}else{xl[i
            +1]=M::op(&xl[i],&dph[u]);}}for i in(0..m).rev(){let(u,_)=g[v][i];if u==par[v]{xr[i]=xr[i+1].clone();}else{xr[i]=M::op(&dph[u],&xr[i+1]
            );}}for i in 0..m{let(u,_)=g[v][i];if u!=par[v]{dph[u]=M::op(&xl[i],&xr[i+1]);}}dp[v]=xr[0].clone();dpl[v]=V::act(&g.vertex(v),&dp[v]
            );for(u,w)in&g[v]{if*u==par[v]{dph[v]=E::act(&w,&dpl[v]);}}}dp[0]=V::act(&g.vertex(0),&dp[0]);for&(u,_)in&g[0]{dph[u]=V::act(&g.vertex(0
            ),&dph[u]);}for&v in&ord{for(u,w)in&g[v]{if*u==par[v]{continue;}let mut x=E::act(&w,&dph[*u]);for&(vv,_)in&g[*u]{if vv==v{continue
            ;}dph[vv]=M::op(&dph[vv],&x);dph[vv]=V::act(&g.vertex(*u),&dph[vv]);}x=M::op(&dp[*u],&x);dp[*u]=V::act(&g.vertex(*u),&x);}}Self{par,dp
            ,dpl,dph}}pub fn prod(&self,v:usize)->T{self.dp[v].clone()}pub fn prod_subtree(&self,v:usize,p:usize)->T{if p==!0{self.dp[v].clone()}else
            if self.par[v]==p{self.dpl[v].clone()}else{self.dph[p].clone()}}pub fn par(&self,v:usize)->usize{self.par[v]}}}
}
pub(crate) mod macros {
pub mod algebraic {pub use crate::{__cargo_equip_macro_def_algebraic_act as act,__cargo_equip_macro_def_algebraic_algebra as algebra
            ,__cargo_equip_macro_def_algebraic_group as group,__cargo_equip_macro_def_algebraic_monoid as monoid};}
pub mod graph {}
pub mod re_rooting_dp {}
}
pub(crate) mod prelude {pub use crate::__cargo_equip::crates::*;}
mod preludes {
pub mod algebraic {}
pub mod graph {}
pub mod re_rooting_dp {pub(in crate::__cargo_equip)use crate::__cargo_equip::crates::{algebraic,graph};}
}
}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0