結果

問題 No.2892 Lime and Karin
ユーザー 37kt37kt
提出日時 2024-12-05 10:51:05
言語 Rust
(1.77.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
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
5,248 KB
testcase_01 AC 1 ms
5,248 KB
testcase_02 AC 1 ms
5,248 KB
testcase_03 AC 1 ms
5,248 KB
testcase_04 AC 1 ms
5,248 KB
testcase_05 AC 1 ms
5,248 KB
testcase_06 AC 1 ms
5,248 KB
testcase_07 AC 1 ms
5,248 KB
testcase_08 AC 1 ms
5,248 KB
testcase_09 AC 1 ms
5,248 KB
testcase_10 AC 1 ms
5,248 KB
testcase_11 AC 1 ms
5,248 KB
testcase_12 AC 1 ms
5,248 KB
testcase_13 AC 1 ms
5,248 KB
testcase_14 AC 203 ms
22,860 KB
testcase_15 AC 181 ms
18,304 KB
testcase_16 AC 302 ms
26,152 KB
testcase_17 AC 72 ms
8,316 KB
testcase_18 AC 301 ms
26,552 KB
testcase_19 AC 419 ms
37,416 KB
testcase_20 AC 24 ms
5,248 KB
testcase_21 AC 229 ms
21,744 KB
testcase_22 AC 293 ms
26,756 KB
testcase_23 AC 121 ms
12,380 KB
testcase_24 AC 427 ms
36,396 KB
testcase_25 AC 425 ms
36,652 KB
testcase_26 AC 427 ms
38,312 KB
testcase_27 AC 434 ms
40,876 KB
testcase_28 AC 430 ms
36,652 KB
testcase_29 AC 429 ms
39,080 KB
testcase_30 AC 419 ms
32,688 KB
testcase_31 AC 419 ms
32,296 KB
testcase_32 AC 439 ms
44,844 KB
testcase_33 AC 428 ms
36,268 KB
testcase_34 AC 253 ms
32,824 KB
testcase_35 AC 252 ms
32,920 KB
testcase_36 AC 307 ms
32,824 KB
testcase_37 AC 309 ms
32,948 KB
testcase_38 AC 309 ms
32,824 KB
testcase_39 AC 305 ms
35,756 KB
testcase_40 AC 302 ms
39,848 KB
testcase_41 AC 365 ms
36,524 KB
testcase_42 AC 365 ms
40,236 KB
testcase_43 AC 362 ms
37,676 KB
testcase_44 AC 148 ms
15,744 KB
testcase_45 AC 354 ms
29,920 KB
testcase_46 AC 176 ms
16,696 KB
testcase_47 AC 239 ms
25,808 KB
testcase_48 AC 99 ms
11,260 KB
testcase_49 AC 270 ms
24,780 KB
testcase_50 AC 327 ms
32,680 KB
testcase_51 AC 324 ms
33,068 KB
testcase_52 AC 361 ms
34,480 KB
testcase_53 AC 364 ms
32,688 KB
testcase_54 AC 365 ms
36,140 KB
権限があれば一括ダウンロードができます

ソースコード

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