結果

問題 No.2618 除霊
ユーザー 37kt37kt
提出日時 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
このコードへのチャレンジ
(要ログイン)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
5,248 KB
testcase_01 AC 2 ms
5,248 KB
testcase_02 AC 1 ms
5,248 KB
testcase_03 AC 1 ms
5,248 KB
testcase_04 AC 2 ms
5,248 KB
testcase_05 AC 1 ms
5,248 KB
testcase_06 AC 399 ms
46,016 KB
testcase_07 AC 390 ms
45,928 KB
testcase_08 AC 385 ms
47,892 KB
testcase_09 AC 487 ms
47,268 KB
testcase_10 AC 388 ms
48,936 KB
testcase_11 AC 403 ms
55,772 KB
testcase_12 AC 487 ms
48,272 KB
testcase_13 AC 401 ms
66,520 KB
testcase_14 AC 420 ms
68,860 KB
testcase_15 AC 378 ms
64,888 KB
testcase_16 AC 362 ms
64,432 KB
testcase_17 AC 398 ms
67,608 KB
testcase_18 AC 424 ms
66,220 KB
testcase_19 AC 378 ms
65,864 KB
testcase_20 AC 406 ms
65,380 KB
testcase_21 AC 395 ms
63,924 KB
testcase_22 AC 408 ms
68,836 KB
testcase_23 AC 417 ms
65,676 KB
testcase_24 AC 378 ms
64,872 KB
testcase_25 AC 439 ms
67,272 KB
testcase_26 AC 376 ms
65,844 KB
testcase_27 AC 377 ms
64,112 KB
testcase_28 AC 382 ms
66,472 KB
testcase_29 AC 420 ms
67,096 KB
testcase_30 AC 380 ms
68,172 KB
testcase_31 AC 427 ms
48,476 KB
testcase_32 AC 392 ms
47,092 KB
testcase_33 AC 368 ms
47,232 KB
testcase_34 AC 443 ms
45,592 KB
testcase_35 AC 384 ms
46,844 KB
testcase_36 AC 433 ms
48,204 KB
testcase_37 AC 407 ms
47,012 KB
testcase_38 AC 369 ms
48,504 KB
testcase_39 AC 383 ms
48,920 KB
testcase_40 AC 366 ms
48,364 KB
testcase_41 AC 426 ms
49,484 KB
testcase_42 AC 512 ms
49,256 KB
権限があれば一括ダウンロードができます

ソースコード

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