結果

問題 No.901 K-ary εxtrεεmε
ユーザー 37kt
提出日時 2025-03-07 16:18:53
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 66 ms / 3,000 ms
コード長 12,547 bytes
コンパイル時間 14,223 ms
コンパイル使用メモリ 401,232 KB
実行使用メモリ 15,488 KB
最終ジャッジ日時 2025-03-07 16:19:12
合計ジャッジ時間 19,360 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 29
権限があれば一括ダウンロードができます

ソースコード

diff # 

// verification-helper: PROBLEM https://yukicoder.me/problems/no/901

pub use __cargo_equip::prelude::*;

use heavy_light_decomposition::HeavyLightDecomposition;
use proconio::{fastout, input};

#[fastout]
fn main() {
    input! {
        n: usize,
        uvw: [(usize, usize, i64); n - 1],
        q: usize,
        vs: [[usize]; q],
    }

    let hld = HeavyLightDecomposition::from_edges(&uvw, 0);
    let mut ws = vec![0; n - 1];
    for &(u, v, w) in &uvw {
        ws[hld.edge_index(u, v)] = w;
    }

    let mut sum_ws = vec![0; n];
    for i in 0..n - 1 {
        sum_ws[i + 1] = sum_ws[i] + ws[i];
    }

    for vs in vs {
        let (_, map) = hld.compress(&vs);
        let mut res = 0;
        for i in 0..map.len() {
            let s = map[i];
            let t = map[(i + 1) % map.len()];
            hld.path_query(s, t, false, |l, r, _| {
                res += sum_ws[r] - sum_ws[l];
            });
        }
        println!("{}", res / 2);
    }
}

// The following code was expanded by `cargo-equip`.

///  # Bundled libraries
/// 
///  - `csr_array 0.1.0 (path+████████████████████████████████████████████████████)`                       published in **missing** licensed under `CC0-1.0` as `crate::__cargo_equip::crates::__csr_array_0_1_0`
///  - `heavy_light_decomposition 0.1.0 (path+██████████████████████████████████████████████████████████)` published in **missing** licensed under `CC0-1.0` as `crate::__cargo_equip::crates::heavy_light_decomposition`
#[cfg_attr(any(), rustfmt::skip)]
#[allow(unused)]
mod __cargo_equip {
    pub(crate) mod crates {
        pub mod __csr_array_0_1_0 {mod builder{use crate::__cargo_equip::crates::__csr_array_0_1_0::CsrArray;#[derive(Clone)]pub struct CsrArrayBuilder<T>{pub(in crate::__cargo_equip::crates::__csr_array_0_1_0)n:usize,pub(in crate::__cargo_equip::crates::__csr_array_0_1_0)idx:Vec<usize>,pub(in crate::__cargo_equip::crates::__csr_array_0_1_0)val:Vec<T>,}impl<T>CsrArrayBuilder<T>{pub fn new(n:usize)->Self{Self{n,idx:vec![],val:vec![],}}pub fn with_capacity(n:usize,capacity:usize)->Self{Self{n,idx:Vec::with_capacity(capacity),val:Vec::with_capacity(capacity),}}pub fn len(&self)->usize{self.n}pub fn is_empty(&self)->bool{self.n==0}pub fn push(&mut self,i:usize,x:T){self.idx.push(i);self.val.push(x);}pub fn build(mut self)->CsrArray<T>{let mut sep=vec![0;self.n+1];for&i in&self.idx{sep[i]+=1;}for i in 0..self.n{sep[i+1]+=sep[i];}let mut ord=vec![0;self.idx.len()];for(j,&i)in self.idx.iter().enumerate().rev(){sep[i]-=1;ord[j]=sep[i];}for i in 0..self.idx.len(){while ord[i]!=i{let j=ord[i];ord.swap(i,j);self.val.swap(i,j);}}CsrArray{sep,val:self.val}}}}mod csr_array{use std::{iter::FusedIterator,ops::Index};use crate::__cargo_equip::crates::__csr_array_0_1_0::CsrArrayBuilder;impl<T>CsrArray<T>{pub fn new(n:usize,idx_val:impl IntoIterator<Item=(usize,T)>)->Self{let(idx,val)=idx_val.into_iter().unzip();let builder=CsrArrayBuilder{n,idx,val};builder.build()}pub fn len(&self)->usize{self.sep.len()-1}pub fn is_empty(&self)->bool{self.sep.len()==1}pub fn flat_len(&self)->usize{self.val.len()}pub fn iter(&self)->Iter<'_,T>{Iter{sep:&self.sep,val:&self.val,}}}pub struct Iter<'a,T>{sep:&'a[usize],val:&'a[T],}impl<'a,T>Iterator for Iter<'a,T>{type Item=&'a[T];fn next(&mut self)->Option<Self::Item>{let&[l,r,..]=self.sep else{return None;};let len=r-l;let(head,tail)=self.val.split_at(len);self.sep=&self.sep[1..];self.val=tail;Some(head)}}impl<'a,T>DoubleEndedIterator for Iter<'a,T>{fn next_back(&mut self)->Option<Self::Item>{let&[..,l,r]=self.sep else{return None;};let len=r-l;let(head,tail)=self.val.split_at(self.val.len()-len);self.sep=&self.sep[..self.sep.len()-1];self.val=head;Some(tail)}}impl<'a,T>ExactSizeIterator for Iter<'a,T>{fn len(&self)->usize{self.sep.len()-1}}impl<'a,T>FusedIterator for Iter<'a,T>{}impl<'a,T>IntoIterator for&'a CsrArray<T>{type Item=&'a[T];type IntoIter=Iter<'a,T>;fn into_iter(self)->Self::IntoIter{self.iter()}}impl<T>Index<usize>for CsrArray<T>{type Output=[T];fn index(&self,index:usize)->&Self::Output{assert!(index<self.len());&self.val[self.sep[index]..self.sep[index+1]]}}impl<T:std::fmt::Debug>std::fmt::Debug for CsrArray<T>{fn fmt(&self,f:&mut std::fmt::Formatter<'_>)->std::fmt::Result{write!(f,"[")?;for(i,v)in self.iter().enumerate(){write!(f,"{:?}",v)?;if i<self.len()-1{write!(f,", ")?;}}write!(f,"]")}}#[derive(Clone)]pub struct CsrArray<T>{pub(in crate::__cargo_equip::crates::__csr_array_0_1_0)sep:Vec<usize>,pub(in crate::__cargo_equip::crates::__csr_array_0_1_0)val:Vec<T>,}}pub use builder::CsrArrayBuilder;pub use csr_array::CsrArray;}
        pub mod heavy_light_decomposition {use crate::__cargo_equip::preludes::heavy_light_decomposition::*;pub mod compress{use crate::__cargo_equip::preludes::heavy_light_decomposition::*;use crate::__cargo_equip::crates::heavy_light_decomposition::HeavyLightDecomposition;impl HeavyLightDecomposition{pub fn compress(&self,vs:&[usize])->(Vec<usize>,Vec<usize>){if vs.is_empty(){return(vec![],vec![]);}let mut v=Vec::with_capacity(vs.len()*2-1);v.extend(vs);v.sort_unstable_by_key(|&v|self.down[v]);for i in 0..v.len()-1{v.push(self.lca(v[i],v[i+1]));}v.sort_unstable_by_key(|&v|self.down[v]);v.dedup();let sz=v.len();let mut par=vec![0;sz];let mut st=vec![0];for i in 1..sz{while!self.in_subtree(*st.last().unwrap(),v[i]){st.pop();}par[i]=*st.last().unwrap();st.push(i);}(par,v)}}}pub mod construct{use crate::__cargo_equip::preludes::heavy_light_decomposition::*;use crate::__cargo_equip::crates::heavy_light_decomposition::{Edge,HeavyLightDecomposition};impl HeavyLightDecomposition{pub fn from_parents(par:&[usize])->Self{let n=par.len();let root=0;let mut down=vec![-1;n];let mut next=std::iter::once(-1).chain(par.iter().skip(1).map(|&p|p as i32)).collect::<Vec<_>>();let mut sub=vec![1;n];for v in(1..n).rev(){let p=next[v]as usize;sub[p]+=sub[v];down[p]=down[p].max(sub[v]);}for v in(1..n).rev(){let p=next[v]as usize;if down[p]==sub[v]{sub[v]=!sub[v];down[p]=!down[p];}}sub[root]=!down[root]+1;down[root]=0;for v in 1..n{let p=next[v]as usize;let nsub=!down[v]+1;if sub[v]<0{down[v]=down[p]+1;next[v]=if next[p]<0{p as i32}else{next[p]};}else{down[v]=down[p]+sub[p];sub[p]+=sub[v];next[v]=!(p as i32);}sub[v]=nsub;}let mut tour=vec![-1;n];for v in 0..n{let t=down[v]as usize;tour[t]=v as i32;}Self{n,root,down,next,sub,tour,}}pub fn from_edges(edges:&[impl Edge],root:usize)->Self{let n=edges.len()+1;assert!(root<n);let mut down=vec![0;n];let mut next=vec![0;n];for e in edges{let(u,v)=e.endpoints();down[u]+=1;down[v]+=1;next[u]^=v as i32;next[v]^=u as i32;}let mut tour=Vec::with_capacity(n);for v in 0..n{if v!=root&&down[v]==1{tour.push(v as i32);}}for i in 0..n-1{let v=tour[i]as usize;down[v]=-1;let u=next[v]as usize;next[u]^=v as i32;down[u]-=1;if down[u]==1&&u!=root{tour.push(u as i32);}}let mut sub=vec![1;n];for&v in&tour{let v=v as usize;let p=next[v]as usize;sub[p]+=sub[v];down[p]=down[p].max(sub[v]);}for&v in&tour{let v=v as usize;let p=next[v]as usize;if down[p]==sub[v]{sub[v]=!sub[v];down[p]=!down[p];}}sub[root]=!down[root]+1;down[root]=0;next[root]=-1;for&v in tour.iter().rev(){let v=v as usize;let p=next[v]as usize;let nsub=!down[v]+1;if sub[v]<0{down[v]=down[p]+1;next[v]=if next[p]<0{p as i32}else{next[p]};}else{down[v]=down[p]+sub[p];sub[p]+=sub[v];next[v]=!(p as i32);}sub[v]=nsub;}tour.resize(n,-1);for v in 0..n{let t=down[v]as usize;tour[t]=v as i32;}Self{n,root,down,next,sub,tour,}}}}use std::iter::FusedIterator;pub struct HeavyLightDecomposition{n:usize,root:usize,down:Vec<i32>,next:Vec<i32>,sub:Vec<i32>,tour:Vec<i32>,}impl HeavyLightDecomposition{pub fn len(&self)->usize{self.n}pub fn root(&self)->usize{self.root}pub fn head(&self,v:usize)->usize{if self.next[v]<0{v}else{self.next[v]as usize}}pub fn la(&self,mut v:usize,d:usize)->Option<usize>{let mut d=d as i32;while v!=!0{let u=self.head(v);if self.down[v]-d>=self.down[u]{v=self.tour[(self.down[v]-d)as usize]as usize;break;}d-=self.down[v]-self.down[u]+1;v=if u==self.root{!0}else{(!self.next[u])as usize};}if v==!0{None}else{Some(v)}}pub fn lca(&self,mut u:usize,mut v:usize)->usize{let mut du=self.down[u];let mut dv=self.down[v];if du>dv{std::mem::swap(&mut u,&mut v);std::mem::swap(&mut du,&mut dv);}if dv<du+self.sub[u]{return u;}while du<dv{v=!self.next[self.head(v)]as usize;dv=self.down[v];}v}pub fn dist(&self,mut u:usize,mut v:usize)->usize{let mut dist=0;while self.head(u)!=self.head(v){if self.down[u]>self.down[v]{std::mem::swap(&mut u,&mut v);}dist+=self.down[v]-self.down[self.head(v)]+1;v=!self.next[self.head(v)]as usize;}dist+=(self.down[u]-self.down[v]).abs();dist as usize}pub fn jump(&self,mut s:usize,mut t:usize,d:usize)->Option<usize>{let(ss,tt)=(s,t);let(mut dist_sl,mut dist_tl)=(0,0);while self.head(s)!=self.head(t){if self.down[s]>self.down[t]{dist_sl+=self.down[s]-self.down[self.head(s)]+1;s=!self.next[self.head(s)]as usize;}else{dist_tl+=self.down[t]-self.down[self.head(t)]+1;t=!self.next[self.head(t)]as usize;}}if self.down[s]>self.down[t]{dist_sl+=self.down[s]-self.down[t];}else{dist_tl+=self.down[t]-self.down[s];}let mut d=d as i32;if d<=dist_sl{return Some(self.la(ss,d as usize).unwrap());}d-=dist_sl;if d<=dist_tl{return Some(self.la(tt,(dist_tl-d)as usize).unwrap());}None}pub fn parent(&self,v:usize)->Option<usize>{if v==self.root{None}else if self.next[v]<0{Some(!self.next[v]as usize)}else{Some(self.tour[(self.down[v]-1)as usize]as usize)}}pub fn vertex_index(&self,v:usize)->usize{self.down[v]as usize}pub fn edge_index(&self,u:usize,v:usize)->usize{debug_assert!(self.parent(u)==Some(v)||self.parent(v)==Some(u));(self.down[u].max(self.down[v])-1)as usize}pub fn subtree_size(&self,v:usize)->usize{self.sub[v]as usize}pub fn subtree_range(&self,v:usize)->(usize,usize){let l=self.down[v]as usize;let r=(self.down[v]+self.sub[v])as usize;(l,r)}pub fn in_subtree(&self,r:usize,v:usize)->bool{let(l1,r1)=self.subtree_range(r);let(l2,r2)=self.subtree_range(v);l1<=l2&&r2<=r1}pub fn dist_table(&self,mut s:usize)->Vec<usize>{let mut dist=vec![!0;self.n];dist[s]=0;while let Some(p)=self.parent(s){dist[p]=dist[s]+1;s=p;}for v in 0..self.n{if dist[v]==!0{dist[v]=dist[self.parent(v).unwrap()]+1;}}dist}pub fn diameter(&self)->(usize,(usize,usize)){let depth=self.dist_table(self.root);let(_,u)=depth.iter().zip(0..).max().unwrap();let from_u=self.dist_table(u);let(_,v)=from_u.iter().zip(0..).max().unwrap();(from_u[v],(u,v))}pub fn path(&self,mut s:usize,mut t:usize)->Vec<usize>{let d=self.dist(s,t);let mut path=vec![!0;d+1];let(mut i,mut j)=(0,d);while s!=t{if self.down[s]>self.down[t]{path[i]=s;i+=1;s=self.parent(s).unwrap();}else{path[j]=t;j-=1;t=self.parent(t).unwrap();}}path[i]=s;path}pub fn path_query(&self,mut s:usize,mut t:usize,vertex_query:bool,mut f:impl FnMut(usize,usize,bool),){let mut f=|l,r,reverse|{if vertex_query{f(l,r,reverse);}else{f(l-1,r-1,reverse);}};let mut down_query=vec![];while self.head(s)!=self.head(t){if self.down[s]<self.down[t]{let l=self.down[self.head(t)];let r=self.down[t]+1;down_query.push((l,r));t=!self.next[self.head(t)]as _;}else{let l=self.down[self.head(s)];let r=self.down[s]+1;f(l as _,r as _,true);s=!self.next[self.head(s)]as _;}}if vertex_query{if self.down[s]<self.down[t]{let l=self.down[s];let r=self.down[t]+1;f(l as _,r as _,false);}else{let l=self.down[t];let r=self.down[s]+1;f(l as _,r as _,true);}}else if self.down[s]<self.down[t]{let l=self.down[s]+1;let r=self.down[t]+1;f(l as _,r as _,false);}else if self.down[s]>self.down[t]{let l=self.down[t]+1;let r=self.down[s]+1;f(l as _,r as _,true);}for&(l,r)in down_query.iter().rev(){f(l as _,r as _,false);}}pub fn euler_tour(&self,)->impl Iterator<Item=usize>+FusedIterator+ExactSizeIterator+DoubleEndedIterator+'_{self.tour.iter().map(|&v|v as usize)}}pub trait Edge{fn endpoints(&self)->(usize,usize);}impl Edge for(usize,usize){fn endpoints(&self)->(usize,usize){*self}}impl<T>Edge for(usize,usize,T){fn endpoints(&self)->(usize,usize){(self.0,self.1)}}}
    }

    pub(crate) mod macros {
        pub mod __csr_array_0_1_0 {}
        pub mod heavy_light_decomposition {}
    }

    pub(crate) mod prelude {pub use crate::__cargo_equip::crates::*;}

    mod preludes {
        pub mod __csr_array_0_1_0 {}
        pub mod heavy_light_decomposition {pub(in crate::__cargo_equip)use crate::__cargo_equip::crates::__csr_array_0_1_0 as csr_array;}
    }
}
0