結果

問題 No.2892 Lime and Karin
ユーザー 37kt37kt
提出日時 2024-12-05 10:51:05
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 439 ms / 8,000 ms
コード長 8,211 bytes
コンパイル時間 20,396 ms
コンパイル使用メモリ 378,612 KB
実行使用メモリ 44,844 KB
最終ジャッジ日時 2024-12-05 10:51:42
合計ジャッジ時間 29,983 ms
ジャッジサーバーID
(参考情報)
judge2 / judge4
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 52
権限があれば一括ダウンロードができます

ソースコード

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

pub use __cargo_equip::prelude::*;
use centroid_decomposition::centroid_decomposition;
use graph::Graph;
#[allow(unused_imports)]
use proconio::{
input,
marker::{Bytes, Chars, Usize1},
};
fn main() {
input! {
n: usize,
uv: [(Usize1, Usize1); n - 1],
s: Bytes,
}
let a = s
.iter()
.map(|&c| if c == b'1' { 1 } else { -1 })
.collect::<Vec<_>>();
let g = Graph::from_unweighted_undirected_edges(n, &uv);
let mut f = vec![0; n];
let mut res = 0;
centroid_decomposition(&g, &mut |idx, par, m| {
let r = par[0];
let mut x = vec![];
let mut y = vec![];
f[r] = 0;
for i in 0..idx.len() {
let v = idx[i];
let p = par[i];
f[v] = f[p] + a[v];
if i < m {
x.push(f[v] + a[r]);
} else {
y.push(f[v]);
}
}
x.sort();
y.sort();
y.reverse();
let mut j = 0;
for i in 0..x.len() {
while j < y.len() && x[i] + y[j] > 0 {
j += 1;
}
res += j;
}
});
for v in 0..n {
if a[v] == 1 {
res += 1;
}
}
for &(u, v) in &uv {
if a[u] == 1 && a[v] == 1 {
res += 1;
}
}
println!("{}", res);
}
// The following code was expanded by `cargo-equip`.
/// # Bundled libraries
///
/// - `centroid-decomposition 0.1.0 (path+███████████████████████████████████████████████████████)` published in **missing** licensed under `CC0-1.0`
    as `crate::__cargo_equip::crates::centroid_decomposition`
/// - `graph 0.1.0 (path+███████████████████████████████████████)` published in **missing** licensed under `CC0-1.0`
    as `crate::__cargo_equip::crates::graph`
#[cfg_attr(any(), rustfmt::skip)]
#[allow(unused)]
mod __cargo_equip {
pub(crate) mod crates {
pub mod centroid_decomposition {use crate::__cargo_equip::preludes::centroid_decomposition::*;use graph::Graph;fn dfs1(v:usize,p:usize,g
            :&Graph<(),()>,sz:&mut[usize]){sz[v]=1;for&(u,_)in&g[v]{if u==p{continue;}dfs1(u,v,g,sz);sz[v]+=sz[u];}}fn dfs2(v:usize,p:usize,mid:usize
            ,g:&Graph<(),()>,sz:&[usize])->usize{for&(u,_)in&g[v]{if u==p{continue;}if sz[u]>mid{return dfs2(u,v,mid,g,sz);}}v}fn dfs4(v:usize,p:usize
            ,g:&Graph<(),()>,pre_idx:&[usize],idx:&mut Vec<usize>,par:&mut Vec<usize>,){idx.push(pre_idx[v]);par.push(pre_idx[p]);for&(u,_)in&g[v]{if
            u==p{continue;}dfs4(u,v,g,pre_idx,idx,par);}}fn dfs3(v:usize,g:&Graph<(),()>,pre_idx:&[usize],conv:&mut[usize],f:&mut impl FnMut(&[usize]
            ,&[usize],usize),){let n=g.len();let mut sz=vec![0;n];dfs1(v,!0,g,&mut sz);let c=dfs2(v,!0,sz[v]/2,g,&sz);dfs1(c,!0,g,&mut sz);let n=sz[c]
            ;if n<=2{return;}let mut szsum=0;let mut rl=vec![];let mut rr=vec![];for&(u,_)in&g[c]{if szsum+sz[u]<=(n-1)/2{szsum+=sz[u];rl.push(u
            );}else{rr.push(u);}}conv[pre_idx[c]]=0;let mut idx_l=vec![];let mut par_l=vec![];let mut es_l=vec![];for&u in&rl{dfs4(u,c,g,pre_idx,&mut
            idx_l,&mut par_l);}for i in 0..idx_l.len(){conv[idx_l[i]]=i+1;es_l.push((conv[par_l[i]],i+1));}let mut idx_r=vec![];let mut par_r=vec![]
            ;let mut es_r=vec![];for&u in&rr{dfs4(u,c,g,pre_idx,&mut idx_r,&mut par_r);}for i in 0..idx_r.len(){conv[idx_r[i]]=i+1;es_r.push
            ((conv[par_r[i]],i+1));}let gl=Graph::from_unweighted_undirected_edges(idx_l.len()+1,&es_l);let gr=Graph::from_unweighted_undirected_edges
            (idx_r.len()+1,&es_r);let mut idx=vec![];idx.append(&mut idx_l.clone());idx.append(&mut idx_r.clone());let mut par=vec![];par.append(&mut
            par_l);par.append(&mut par_r);f(&idx,&par,idx_l.len());idx_l.insert(0,pre_idx[c]);idx_r.insert(0,pre_idx[c]);dfs3(0,&gl,&idx_l,conv,f
            );dfs3(0,&gr,&idx_r,conv,f);}pub fn centroid_decomposition(g:&Graph<(),()>,f:&mut impl FnMut(&[usize],&[usize],usize)){let n=g.len();let
            mut conv=vec![!0;n];dfs3(0,g,&(0..n).collect::<Vec<_>>(),&mut conv,f);}}
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(crate) mod macros {
pub mod centroid_decomposition {}
pub mod graph {}
}
pub(crate) mod prelude {pub use crate::__cargo_equip::crates::*;}
mod preludes {
pub mod centroid_decomposition {pub(in crate::__cargo_equip)use crate::__cargo_equip::crates::graph;}
pub mod graph {}
}
}
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
0