#![allow(unused_imports)] use std::io::{Read, Write}; #[allow(unused_macros)] macro_rules! get { ( $in:ident, [$a:tt; $num:expr] ) => { { let n = $num; (0 .. n).map(|_| get!($in, $a)).collect::>() } }; ( $in:ident, ($($type:ty),*) ) => { ($(get!($in, $type)),*) }; ( $in:ident, $type:ty ) => { { let token = $in.next().unwrap(); token.parse::<$type>().expect( format!("cannot convert \"{}\" into {}", token, stringify!($type)).as_str() ) } }; } macro_rules! input { ( @inner $in:ident, mut $name:ident : $type:tt ) => { let mut $name = get!($in, $type); }; ( @inner $in:ident, $name:ident : $type:tt ) => { let $name = get!($in, $type); }; ( $in:ident, $($($names:ident)* : $type:tt),* ) => { $( input!(@inner $in, $($names)* : $type); )* } } #[allow(unused_macros)] macro_rules! io { ( $in:ident, $out:ident ) => { let mut s = String::new(); std::io::stdin().read_to_string(&mut s).unwrap(); let mut $in = s.split_ascii_whitespace(); let $out = std::io::stdout(); let mut $out = std::io::BufWriter::new($out.lock()); }; } pub mod main { use super::*; use haar_lib::{ algebra::update_sum::*, ds::lazy_segtree::*, //chmin, chmax, //mul_vec, tree::{depth_query::*, *}, }; #[derive(Clone, Default)] pub struct Problem {/* write variables here */} impl Problem { pub fn main(&mut self) -> Result<(), Box> { io!(cin, cout); input!(cin, n: usize); let mut tree = Tree::new(n); for _ in 0..n - 1 { input!(cin, u: usize, v: usize); tree.add_undirected(vec![(u, v, ())]); } let res = TreeDepthQuery::new(&tree, 0); input!(cin, a: [u64; n]); let mut seg = LazySegmentTree::new(n, UpdateSum::::new()); for i in 0..n { let (l, r) = res.me_query(i); seg.update(l..r, Some(a[i])); } input!(cin, q: usize); for _ in 0..q { input!(cin, x: usize); let mut ans = 0; let mut f = |(l, r)| { ans += seg.fold(l..r); seg.update(l..r, Some(0)); }; if let Some(x) = res.ancestor(x, 2) { f(res.me_query(x)); } if let Some(x) = res.ancestor(x, 1) { f(res.me_query(x)); res.children_query(x, 1).into_iter().for_each(|p| f(p)); } f(res.me_query(x)); res.children_query(x, 1).into_iter().for_each(|p| f(p)); res.children_query(x, 2).into_iter().for_each(|p| f(p)); let (l, r) = res.me_query(x); seg.update(l..r, Some(ans)); writeln!(cout, "{}", ans)?; } Ok(()) } /* write functions here */ } } fn main() { main::Problem::default().main().unwrap(); } use crate as haar_lib; pub mod algebra { pub mod action { pub trait Action { type FType; type UType; fn fold_id(&self) -> Self::FType; fn fold(&self, x: Self::FType, y: Self::FType) -> Self::FType; fn update_id(&self) -> Self::UType; fn update(&self, next: Self::UType, cur: Self::UType) -> Self::UType; fn convert(&self, x: Self::FType, y: Self::UType, l: usize) -> Self::FType; } } pub mod update_sum { use crate::algebra::action::Action; use std::{ marker::PhantomData, ops::{Add, Mul}, }; #[derive(Clone, Default)] pub struct UpdateSum(PhantomData, PhantomData); impl UpdateSum { pub fn new() -> Self { Self(PhantomData, PhantomData) } } impl Action for UpdateSum where T: Add + Default + From, U: Mul + Default + From, { type FType = T; type UType = Option; fn fold_id(&self) -> Self::FType { T::default() } fn fold(&self, x: Self::FType, y: Self::FType) -> Self::FType { x + y } fn update_id(&self) -> Self::UType { None } fn update(&self, x: Self::UType, y: Self::UType) -> Self::UType { match x { Some(_) => x, _ => y, } } fn convert(&self, x: Self::FType, y: Self::UType, l: usize) -> Self::FType { match y { Some(y) => T::from(y * U::from(l as u64)), _ => x, } } } } } pub mod algo { pub mod bsearch { #[doc = " x以上となる最小のindexを求める。"] #[doc = ""] #[doc = " # Complexity"] #[doc = " Time complexity $O(\\log(n))$"] pub fn lower_bound(a: &[T], value: &T) -> usize { let n = a.len(); let mut lb = 0; let mut len = n; while len > 0 { let half = len / 2; let mid = lb + half; if &a[mid] < value { len -= half + 1; lb = mid + 1; } else { len = half; } } lb } #[doc = " xを超える最小のindexを求める。"] #[doc = ""] #[doc = " # Complexity"] #[doc = " Time complexity $O(\\log(n))$"] pub fn upper_bound(a: &[T], value: &T) -> usize { let n = a.len(); let mut ub = 0; let mut len = n; while len > 0 { let half = len / 2; let mid = ub + half; if &a[mid] <= value { len -= half + 1; ub = mid + 1; } else { len = half; } } ub } #[doc = " lower_bound, upper_boundの組を求める。"] #[doc = ""] #[doc = " # Complexity"] #[doc = " Time complexity $O(\\log(n))$"] pub fn equal_range(a: &[T], value: &T) -> (usize, usize) { (lower_bound(a, value), upper_bound(a, value)) } } } pub mod ds { pub mod traits { pub trait Foldable { type Output; fn fold(&self, range: Idx) -> Self::Output; } pub trait FoldableMut { type Output; fn fold(&mut self, range: Idx) -> Self::Output; } pub trait Assignable { type Value; fn assign(&mut self, i: Idx, value: Self::Value); } pub trait Updatable { type Value; fn update(&mut self, i: Idx, value: Self::Value); } pub trait IndexableMut { type Output; fn get(&mut self, i: Idx) -> Self::Output; } } pub mod lazy_segtree { use crate::algebra::action::Action; pub use crate::ds::traits::*; use std::ops::Range; pub struct LazySegmentTree { size: usize, data: Vec, lazy: Vec, action: A, } impl> LazySegmentTree { pub fn new(n: usize, a: A) -> Self { let size = n.next_power_of_two() * 2; Self { size, data: vec![a.fold_id(); size], lazy: vec![a.update_id(); size], action: a, } } fn propagate(&mut self, i: usize) { if self.lazy[i] == self.action.update_id() { return; } if i < self.size / 2 { self.lazy[i << 1] = self .action .update(self.lazy[i].clone(), self.lazy[i << 1].clone()); self.lazy[i << 1 | 1] = self .action .update(self.lazy[i].clone(), self.lazy[i << 1 | 1].clone()); } let len = (self.size / 2) >> (31 - (i as u32).leading_zeros()); self.data[i] = self .action .convert(self.data[i].clone(), self.lazy[i].clone(), len); self.lazy[i] = self.action.update_id(); } fn propagate_top_down(&mut self, mut i: usize) { let mut temp = vec![]; while i > 1 { i >>= 1; temp.push(i); } for &i in temp.iter().rev() { self.propagate(i); } } fn bottom_up(&mut self, mut i: usize) { while i > 1 { i >>= 1; self.propagate(i << 1); self.propagate(i << 1 | 1); self.data[i] = self .action .fold(self.data[i << 1].clone(), self.data[i << 1 | 1].clone()); } } } impl> Updatable> for LazySegmentTree { type Value = U; fn update(&mut self, Range { start: l, end: r }: Range, x: U) { self.propagate_top_down(l + self.size / 2); if r < self.size / 2 { self.propagate_top_down(r + self.size / 2); } { let mut l = l + self.size / 2; let mut r = r + self.size / 2; while l < r { if r & 1 == 1 { r -= 1; self.lazy[r] = self.action.update(x.clone(), self.lazy[r].clone()); self.propagate(r); } if l & 1 == 1 { self.lazy[l] = self.action.update(x.clone(), self.lazy[l].clone()); self.propagate(l); l += 1; } r >>= 1; l >>= 1; } } self.bottom_up(l + self.size / 2); if r < self.size / 2 { self.bottom_up(r + self.size / 2); } } } impl> FoldableMut> for LazySegmentTree { type Output = T; fn fold(&mut self, Range { start: l, end: r }: Range) -> Self::Output { self.propagate_top_down(l + self.size / 2); if r < self.size / 2 { self.propagate_top_down(r + self.size / 2); } let mut ret_l = self.action.fold_id(); let mut ret_r = self.action.fold_id(); let mut l = l + self.size / 2; let mut r = r + self.size / 2; while l < r { if r & 1 == 1 { r -= 1; self.propagate(r); ret_r = self.action.fold(self.data[r].clone(), ret_r.clone()); } if l & 1 == 1 { self.propagate(l); ret_l = self.action.fold(ret_l.clone(), self.data[l].clone()); l += 1; } r >>= 1; l >>= 1; } self.action.fold(ret_l, ret_r) } } } } pub mod tree { pub mod depth_query { use crate::{algo::bsearch::lower_bound, tree::*}; use std::collections::VecDeque; pub struct TreeDepthQuery { par: Vec>, depth: Vec, left: Vec, right: Vec, bfs_ord: Vec>, dfs_ord: Vec>, ord: Vec, } impl TreeDepthQuery { pub fn new(tree: &Tree, root: usize) -> Self { let size = tree.len(); let mut ret = Self { par: vec![None; size], depth: vec![0; size], left: vec![0; size], right: vec![0; size], bfs_ord: vec![], dfs_ord: vec![], ord: vec![0; size], }; ret.dfs(&tree, root, None, 0, &mut 0); let mut q = VecDeque::new(); q.push_back((root, 0)); let mut ord = 0; while let Some((i, d)) = q.pop_front() { if ret.bfs_ord.len() <= d { ret.bfs_ord.push(vec![]); } ret.bfs_ord[d].push(ord); ret.ord[i] = ord; ord += 1; for &TreeEdge { to, .. } in tree.nodes[i].neighbors() { if Some(to) != ret.par[i] { q.push_back((to, d + 1)); } } } ret } fn dfs( &mut self, tree: &Tree, cur: usize, par: Option, d: usize, ord: &mut usize, ) { self.par[cur] = par; self.depth[cur] = d; if self.dfs_ord.len() <= d { self.dfs_ord.push(vec![]); } self.dfs_ord[d].push(*ord); self.left[cur] = *ord; *ord += 1; for &TreeEdge { to, .. } in tree.nodes[cur].neighbors() { if Some(to) != par { self.dfs(tree, to, Some(cur), d + 1, ord); } } self.right[cur] = *ord; } pub fn children_query(&self, i: usize, d: usize) -> Option<(usize, usize)> { let d = d + self.depth[i]; if self.bfs_ord.len() > d { let l = lower_bound(&self.dfs_ord[d], &self.left[i]); let r = lower_bound(&self.dfs_ord[d], &self.right[i]); if l >= self.bfs_ord[d].len() { return None; } if r == 0 { return None; } Some((self.bfs_ord[d][l], self.bfs_ord[d][r - 1] + 1)) } else { None } } pub fn me_query(&self, i: usize) -> (usize, usize) { (self.ord[i], self.ord[i] + 1) } pub fn ancestor(&self, i: usize, k: usize) -> Option { let mut p = i; for _ in 0..k { match self.par[p] { Some(x) => p = x, _ => return None, } } Some(p) } } } #[derive(Clone, Debug)] pub struct TreeEdge { pub to: usize, pub weight: T, } #[derive(Clone, Debug)] pub struct TreeNode { pub index: usize, pub parent: Option>, pub children: Vec>, } #[derive(Clone, Debug)] pub struct Tree { pub nodes: Vec>, } impl TreeNode { pub fn neighbors(&self) -> impl DoubleEndedIterator> { self.children.iter().chain(self.parent.iter()) } pub fn neighbors_size(&self) -> usize { self.children.len() + self.parent.as_ref().map_or(0, |_| 1) } } impl Tree { pub fn new(size: usize) -> Self { Self { nodes: (0..size) .map(|i| TreeNode { index: i, parent: None, children: vec![], }) .collect(), } } pub fn add_undirected(&mut self, edges: impl IntoIterator) { for (u, v, w) in edges { self.nodes[u].children.push(TreeEdge { to: v, weight: w }); self.nodes[v].children.push(TreeEdge { to: u, weight: w }); } } pub fn add_directed(&mut self, edges: impl IntoIterator) { for (p, c, w) in edges { assert!(self.nodes[c].parent.is_none()); self.nodes[p].children.push(TreeEdge { to: c, weight: w }); self.nodes[c].parent = Some(TreeEdge { to: p, weight: w }); } } } impl Tree { pub fn len(&self) -> usize { self.nodes.len() } pub fn is_empty(&self) -> bool { self.nodes.is_empty() } } }