use crate::flat_graph::Graph; use proconio::input; fn main() { input! { n: usize, cap: usize, values: [u64; n], edges: [(usize, usize, usize); n - 1], } let cap = cap / 2; let mut g = Graph::from_undirected_edges_with_weight(n, &edges); let (_sorted, _parent) = g.sort_undirected_tree_with_weight(0); let dp = dfs(0, &g, &values, vec![Some(0); cap + 1]); let ans = dp.last().unwrap().unwrap(); println!("{ans}"); } fn dfs( x: usize, g: &Graph<(usize, usize)>, values: &[u64], mut dp: Vec>, ) -> Vec> { for dp in dp.iter_mut().flatten() { *dp += values[x]; } for &(y, w) in &g[x] { if dp.len() <= w { continue; } let mut ep = dp.to_vec(); ep.rotate_right(w); ep[..w].fill(None); let ep = dfs(y, g, values, ep); for (dp, ep) in dp.iter_mut().zip(ep) { *dp = (*dp).max(ep); } } dp } // flat_graph {{{ // https://ngtkana.github.io/ac-adapter-rs/flat_graph/index.html #[allow(unused_imports)] #[allow(dead_code)] mod flat_graph { use std::ops::Index; #[derive(Clone, Debug)] pub struct Graph { start: Vec, tar: Vec, } impl Graph { pub fn sort_undirected_tree(&mut self, root: usize) -> (Vec, Vec) { let n = self.start.len() - 1; assert_eq!(self.tar.len(), 2 * (n - 1)); let mut sorted = vec![]; let mut stack = vec![root]; let mut parent = vec![usize::MAX; n]; parent[root] = 0; while let Some(x) = stack.pop() { sorted.push(x); for &y in &self[x] { if parent[y] != usize::MAX { continue; } parent[y] = x; stack.push(y); } } let mut i = 0; let mut j = 0; for x in 0..n { while j < self.start[x + 1] { if self.tar[j] != parent[x] { self.tar.swap(i, j); i += 1; } j += 1; } self.start[x + 1] = i; } assert_eq!(i, n - 1); self.tar.truncate(n - 1); (sorted, parent) } pub fn from_directed_edges(n: usize, edges: &[(usize, usize)]) -> Self { Self::from_edges_generic( n, edges.len(), edges.iter().map(|&(i, _)| i), edges.iter().map(|&(i, j)| (i, j)), ) } pub fn from_undirected_edges(n: usize, edges: &[(usize, usize)]) -> Self { Self::from_edges_generic( n, edges.len() * 2, edges.iter().flat_map(|&(i, j)| [i, j]), edges.iter().flat_map(|&(i, j)| [(i, j), (j, i)]), ) } } impl Graph<(usize, T)> { pub fn from_directed_edges_with_weight(n: usize, edges: &[(usize, usize, T)]) -> Self { Self::from_edges_generic( n, edges.len(), edges.iter().map(|&(i, _, _)| i), edges.iter().map(|&(i, j, w)| (i, (j, w))), ) } pub fn from_undirected_edges_with_weight(n: usize, edges: &[(usize, usize, T)]) -> Self { Self::from_edges_generic( n, edges.len() * 2, edges.iter().flat_map(|&(i, j, _)| [i, j]), edges .iter() .flat_map(|&(i, j, w)| [(i, (j, w)), (j, (i, w))]), ) } pub fn sort_undirected_tree_with_weight( &mut self, root: usize, ) -> (Vec, Vec) { let n = self.start.len() - 1; assert_eq!(self.tar.len(), 2 * (n - 1)); let mut sorted = vec![]; let mut stack = vec![root]; let mut parent = vec![usize::MAX; n]; parent[root] = 0; while let Some(x) = stack.pop() { sorted.push(x); for &(y, _) in &self[x] { if parent[y] != usize::MAX { continue; } parent[y] = x; stack.push(y); } } let mut i = 0; let mut j = 0; for x in 0..n { while j < self.start[x + 1] { if self.tar[j].0 != parent[x] { self.tar.swap(i, j); i += 1; } j += 1; } self.start[x + 1] = i; } assert_eq!(i, n - 1); self.tar.truncate(n - 1); (sorted, parent) } } impl Graph { fn from_edges_generic( n: usize, m: usize, src: impl Iterator, edges: impl Iterator, ) -> Self { let mut start = vec![0; n + 1]; for i in src { start[i + 1] += 1; } for i in 0..n { start[i + 1] += start[i]; } let edge_count = m; let mut tar = vec![E::default(); edge_count]; for (i, e) in edges { tar[start[i]] = e; start[i] += 1; } start.rotate_right(1); start[0] = 0; Self { start, tar } } } impl Graph { pub fn iter(&self) -> Iter<'_, E> { Iter { index: 0, graph: self, } } } impl Index for Graph { type Output = [E]; fn index(&self, index: usize) -> &Self::Output { &self.tar[self.start[index]..self.start[index + 1]] } } impl<'a, E> IntoIterator for &'a Graph { type Item = &'a [E]; type IntoIter = Iter<'a, E>; fn into_iter(self) -> Self::IntoIter { self.iter() } } pub struct Iter<'a, E> { index: usize, graph: &'a Graph, } impl<'a, E> Iterator for Iter<'a, E> { type Item = &'a [E]; fn next(&mut self) -> Option { if self.index + 1 == self.graph.start.len() { None } else { self.index += 1; Some(&self.graph[self.index - 1]) } } } } // }}}