diff --git a/src/problem/mod.rs b/src/problem/mod.rs index b4b7361..11fa6a4 100644 --- a/src/problem/mod.rs +++ b/src/problem/mod.rs @@ -448,3 +448,5 @@ mod p2239_find_closest_number_to_zero; mod p2218_maximum_value_of_k_coins_from_piles; mod p1561_maximum_number_of_coins_you_can_get; + +mod p2920_maximum_points_after_collecting_coins_from_all_nodes; diff --git a/src/problem/p2920_maximum_points_after_collecting_coins_from_all_nodes.rs b/src/problem/p2920_maximum_points_after_collecting_coins_from_all_nodes.rs new file mode 100644 index 0000000..5aabd2c --- /dev/null +++ b/src/problem/p2920_maximum_points_after_collecting_coins_from_all_nodes.rs @@ -0,0 +1,84 @@ +/** + * [2920] Maximum Points After Collecting Coins From All Nodes + */ +pub struct Solution {} + +// submission codes start here +use std::cell::RefCell; + +struct Search { + memory: RefCell>>, + children: Vec>, + coins: Vec, +} + +impl Search { + fn new(children: Vec>, n: usize, coins: Vec) -> Self { + Self { + children, + // 14是树可能的最大高度 + memory: RefCell::new(vec![vec![-1; 14]; n]), + coins, + } + } + + fn search(&self, node: usize, parent: usize, f: usize, k: i32) -> i32 { + if self.memory.borrow()[node][f] >= 0 { + return self.memory.borrow()[node][f]; + } + + let mut result0 = (self.coins[node] >> f) - k; + let mut result1 = self.coins[node] >> (f + 1); + + for &child in self.children[node].iter() { + if child == parent { + continue; + } + + result0 += self.search(child, node, f, k); + if f + 1 < 14 { + result1 += self.search(child, node, f + 1, k); + } + } + + self.memory.borrow_mut()[node][f] = result0.max(result1); + self.memory.borrow()[node][f] + } +} + +impl Solution { + pub fn maximum_points(edges: Vec>, coins: Vec, k: i32) -> i32 { + let n = coins.len(); + let mut children = vec![vec![]; n]; + + for edge in edges { + let (x, y) = (edge[0] as usize, edge[1] as usize); + + children[x].push(y); + children[y].push(x); + } + + let search = Search::new(children, n, coins); + + search.search(0, usize::MAX, 0, k) + } +} + +// submission codes end + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_2920() { + assert_eq!( + 11, + Solution::maximum_points( + vec![vec![0, 1], vec![1, 2], vec![2, 3]], + vec![10, 10, 3, 3], + 5 + ) + ); + } +}