From 50155415a8952da4064f01979ffbb2f410b60544 Mon Sep 17 00:00:00 2001
From: EFanZh <efanzh@gmail.com>
Date: Sat, 22 May 2021 10:42:31 +0800
Subject: [PATCH] Add problem 623: Add One Row to Tree

---
 src/lib.rs                                    |  1 +
 src/problem_0623_add_one_row_to_tree/mod.rs   | 46 +++++++++++++++
 .../recursive.rs                              | 58 +++++++++++++++++++
 3 files changed, 105 insertions(+)
 create mode 100644 src/problem_0623_add_one_row_to_tree/mod.rs
 create mode 100644 src/problem_0623_add_one_row_to_tree/recursive.rs

diff --git a/src/lib.rs b/src/lib.rs
index 28e59371..491af8f3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -445,6 +445,7 @@ pub mod problem_0606_construct_string_from_binary_tree;
 pub mod problem_0609_find_duplicate_file_in_system;
 pub mod problem_0617_merge_two_binary_trees;
 pub mod problem_0621_task_scheduler;
+pub mod problem_0623_add_one_row_to_tree;
 pub mod problem_0628_maximum_product_of_three_numbers;
 pub mod problem_0647_palindromic_substrings;
 pub mod problem_0657_robot_return_to_origin;
diff --git a/src/problem_0623_add_one_row_to_tree/mod.rs b/src/problem_0623_add_one_row_to_tree/mod.rs
new file mode 100644
index 00000000..5654bdb5
--- /dev/null
+++ b/src/problem_0623_add_one_row_to_tree/mod.rs
@@ -0,0 +1,46 @@
+use super::data_structures::TreeNode;
+use std::cell::RefCell;
+use std::rc::Rc;
+
+pub mod recursive;
+
+pub trait Solution {
+    fn add_one_row(root: Option<Rc<RefCell<TreeNode>>>, val: i32, depth: i32) -> Option<Rc<RefCell<TreeNode>>>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::super::test_utilities;
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            (
+                (&[Some(4), Some(2), Some(6), Some(3), Some(1), Some(5)] as &[_], 1, 2),
+                &[
+                    Some(4),
+                    Some(1),
+                    Some(1),
+                    Some(2),
+                    None,
+                    None,
+                    Some(6),
+                    Some(3),
+                    Some(1),
+                    Some(5),
+                ] as &[_],
+            ),
+            (
+                (&[Some(4), Some(2), None, Some(3), Some(1)], 1, 3),
+                &[Some(4), Some(2), None, Some(1), Some(1), Some(3), None, None, Some(1)],
+            ),
+        ];
+
+        for ((root, val, depth), expected) in test_cases.iter().copied() {
+            assert_eq!(
+                S::add_one_row(test_utilities::make_tree(root.iter().copied()), val, depth),
+                test_utilities::make_tree(expected.iter().copied())
+            );
+        }
+    }
+}
diff --git a/src/problem_0623_add_one_row_to_tree/recursive.rs b/src/problem_0623_add_one_row_to_tree/recursive.rs
new file mode 100644
index 00000000..5a2748c9
--- /dev/null
+++ b/src/problem_0623_add_one_row_to_tree/recursive.rs
@@ -0,0 +1,58 @@
+use super::super::data_structures::TreeNode;
+
+pub struct Solution;
+
+use std::cell::RefCell;
+use std::rc::Rc;
+
+impl Solution {
+    fn new_node_with_left_child(val: i32, node: Option<Rc<RefCell<TreeNode>>>) -> TreeNode {
+        TreeNode {
+            val,
+            left: node,
+            right: None,
+        }
+    }
+
+    fn new_node_with_right_child(val: i32, node: Option<Rc<RefCell<TreeNode>>>) -> TreeNode {
+        TreeNode {
+            val,
+            left: None,
+            right: node,
+        }
+    }
+
+    fn helper(
+        node: &mut Option<Rc<RefCell<TreeNode>>>,
+        val: i32,
+        depth: i32,
+        make_child: impl FnOnce(i32, Option<Rc<RefCell<TreeNode>>>) -> TreeNode,
+    ) {
+        if depth == 1 {
+            *node = Some(Rc::new(RefCell::new(make_child(val, node.take()))));
+        } else if let Some(mut node) = node.as_deref().map(RefCell::borrow_mut) {
+            Self::helper(&mut node.left, val, depth - 1, Self::new_node_with_left_child);
+            Self::helper(&mut node.right, val, depth - 1, Self::new_node_with_right_child);
+        }
+    }
+
+    pub fn add_one_row(mut root: Option<Rc<RefCell<TreeNode>>>, val: i32, depth: i32) -> Option<Rc<RefCell<TreeNode>>> {
+        Self::helper(&mut root, val, depth, Self::new_node_with_left_child);
+
+        root
+    }
+}
+
+impl super::Solution for Solution {
+    fn add_one_row(root: Option<Rc<RefCell<TreeNode>>>, val: i32, depth: i32) -> Option<Rc<RefCell<TreeNode>>> {
+        Self::add_one_row(root, val, depth)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}