結果
問題 | No.2892 Lime and Karin |
ユーザー |
|
提出日時 | 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 |
ソースコード
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]{ifu==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,&mutidx_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(&mutpar_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();letmut conv=vec![!0;n];dfs3(0,g,&(0..n).collect::<Vec<_>>(),&mut conv,f);}}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,{pubfn 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,_,_)inedges{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(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 {}}}