use proconio::input; use std::collections::HashMap; fn main() { input! { n: usize, m: usize, values: [u64; n], edges: [(usize, usize, usize); n - 1], } let mut g = vec![Vec::new(); n]; let mut weight = HashMap::new(); for &(i, j, w) in &edges { g[i].push(j); g[j].push(i); weight.insert((i, j), w); weight.insert((j, i), w); } let mut sorted = Vec::new(); let mut parent = vec![usize::MAX; n]; let mut stack = vec![0]; while let Some(i) = stack.pop() { sorted.push(i); g[i].retain(|&j| j != parent[i]); for &j in &g[i] { stack.push(j); parent[j] = i; } } let weights = (0..n) .map(|i| if i == 0 { 0 } else { weight[&(parent[i], i)] }) .collect::>(); let dp = tree_dp( &O { max: m / 2, weights: &weights, values: &values, }, &g, &sorted, ); let ans = dp[0].f[m / 2]; println!("{ans}"); } #[derive(Debug, Clone, PartialEq, Default)] struct Value { f: Vec, } struct O<'a> { max: usize, weights: &'a [usize], values: &'a [u64], } impl Op for O<'_> { type Acc = Value; type Value = Value; fn f(&self, value: &Self::Value, index: usize) -> Self::Acc { let w = self.weights[index].min(self.max + 1); let mut f = value.f.clone(); f.rotate_right(w); f[..w].iter_mut().for_each(|x| *x = 0); Self::Acc { f } } fn g(&self, acc: &Self::Acc, index: usize) -> Self::Value { let v = self.values[index]; let mut f = acc.f.clone(); f.iter_mut().for_each(|x| *x = *x + v); Self::Value { f } } fn identity(&self) -> Self::Acc { Self::Acc { f: vec![0; self.max + 1], } } fn mul(&self, lhs: &Self::Acc, rhs: &Self::Acc) -> Self::Acc { let mut f = vec![0; self.max + 1]; for (i, &x) in lhs.f.iter().enumerate() { for (j, &y) in rhs.f[..=self.max - i].iter().enumerate() { f[i + j] = f[i + j].max(x + y); } } Self::Acc { f } } } pub trait Op { type Value: Default + Clone; type Acc: Default + Clone; fn f(&self, value: &Self::Value, index: usize) -> Self::Acc; fn g(&self, acc: &Self::Acc, index: usize) -> Self::Value; fn identity(&self) -> Self::Acc; fn mul(&self, lhs: &Self::Acc, rhs: &Self::Acc) -> Self::Acc; } pub fn tree_dp(o: &O, g: &[Vec], sorted: &[usize]) -> Vec { let mut dp = vec![O::Value::default(); g.len()]; for &i in sorted.iter().rev() { dp[i] = o.g( &g[i] .iter() .fold(o.identity(), |acc, &j| o.mul(&acc, &o.f(&dp[j], j))), i, ); } dp }