diff --git a/src/problem/mod.rs b/src/problem/mod.rs index c4d82b4..f8a6f6d 100644 --- a/src/problem/mod.rs +++ b/src/problem/mod.rs @@ -39,4 +39,5 @@ mod p1686_stone_game_vi; mod p1690_stone_game_vii; mod p292_nim_game; mod p1696_jump_game_vi; -mod lcp30_magic_tower; \ No newline at end of file +mod lcp30_magic_tower; +mod p2641_cousins_in_binary_tree_ii; \ No newline at end of file diff --git a/src/problem/p2641_cousins_in_binary_tree_ii.rs b/src/problem/p2641_cousins_in_binary_tree_ii.rs new file mode 100644 index 0000000..ad3ca3a --- /dev/null +++ b/src/problem/p2641_cousins_in_binary_tree_ii.rs @@ -0,0 +1,97 @@ +/** + * [2641] Cousins in Binary Tree II + */ +pub struct Solution {} + +use crate::util::tree::{TreeNode, to_tree}; + +// submission codes start here + +// Definition for a binary tree node. +// #[derive(Debug, PartialEq, Eq)] +// pub struct TreeNode { +// pub val: i32, +// pub left: Option>>, +// pub right: Option>>, +// } +// +// impl TreeNode { +// #[inline] +// pub fn new(val: i32) -> Self { +// TreeNode { +// val, +// left: None, +// right: None +// } +// } +// } +use std::rc::Rc; +use std::cell::RefCell; +impl Solution { + pub fn replace_value_in_tree(root: Option>>) -> Option>> { + let root = if let Some(r) = root { + r + } else { + return None; + }; + + let mut queue = Vec::new(); + queue.push(Rc::clone(&root)); + + while !queue.is_empty() { + let mut next_queue = Vec::new(); + + let mut sum = 0; + for node in &queue { + if let Some(left) = &node.borrow().left { + next_queue.push(Rc::clone(left)); + sum += left.borrow().val; + } + + if let Some(right) = &node.borrow().right { + next_queue.push(Rc::clone(right)); + sum += right.borrow().val; + } + } + + for node in &queue { + let mut children = 0; + if let Some(left) = &node.borrow().left { + children += left.borrow().val; + } + + if let Some(right) = &node.borrow().right { + children += right.borrow().val; + } + + if let Some(left) = &node.borrow().left { + left.borrow_mut().val = sum - children; + } + + if let Some(right) = &node.borrow().right { + right.borrow_mut().val = sum - children; + } + } + + queue = next_queue; + } + + + root.borrow_mut().val = 0; + Some(root) + } +} + +// submission codes end + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_2641() { + assert_eq!(Solution::replace_value_in_tree(to_tree( + vec![Some(5),Some(4),Some(9),Some(1),Some(10),None,Some(7)])), + to_tree(vec![Some(0),Some(0),Some(0),Some(7),Some(7),None,Some(11)])); + } +}