結果
問題 | No.2618 除霊 |
ユーザー |
|
提出日時 | 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 |
ソースコード
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_rules!__cargo_equip_macro_def_algebraic_algebra{($ident:ident,$ty:ty)=>{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_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;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_rules!__cargo_equip_macro_def_algebraic_monoid{($ident:ident,$e:expr,$op:expr)=>{impl$crate::__cargo_equip::crates::algebraic::Monoid for$ident{fn e()->Self::S{$e}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_rules!__cargo_equip_macro_def_algebraic_group{($ident:ident,$e:expr,$op:expr,$inv:expr)=>{impl$crate::__cargo_equip::crates::algebraic::Monoid for$ident{fn e()->Self::S{$e}fn op(x:&Self::S,y:&Self::S)->Self::S{$op(x,y)}}impl$crate::__cargo_equip::crates::algebraic::Group for$ident{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 i32i64 i128);}pub mod graph {use std::ops::Index;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 constGRID_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{ifedges.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,{pubfn from_directed_edges(n:usize,edges:&[(usize,usize,E)])->Self{Self::from_vertices_and_directed_edges(&vec![();n],edges)}pub fnfrom_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 fnfrom_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 fnfrom_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 structReRootingDP<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()}elseif 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};}}}