rustgym/leetcode/
_450_delete_node_in_a_bst.rs1struct Solution;
2use rustgym_util::*;
3use std::cmp::Ordering::*;
4
5trait Delete {
6 fn delete(self, key: i32) -> TreeLink;
7 fn smallest(&self) -> i32;
8}
9
10impl Delete for TreeLink {
11 fn delete(self, key: i32) -> TreeLink {
12 if let Some(link) = self {
13 let mut node = link.borrow_mut();
14 let left = node.left.take();
15 let right = node.right.take();
16 let val = node.val;
17 match key.cmp(&val) {
18 Equal => match (left, right) {
19 (None, None) => None,
20 (Some(left), None) => Some(left),
21 (None, Some(right)) => Some(right),
22 (left, right) => {
23 let smallest = right.smallest();
24 tree!(smallest, left, right.delete(smallest))
25 }
26 },
27 Less => tree!(val, left.delete(key), right),
28 Greater => tree!(val, left, right.delete(key)),
29 }
30 } else {
31 None
32 }
33 }
34 fn smallest(&self) -> i32 {
35 if let Some(node) = self {
36 let node = node.borrow();
37 let val = node.val;
38 if node.left.is_some() {
39 node.left.smallest()
40 } else {
41 val
42 }
43 } else {
44 unreachable!()
45 }
46 }
47}
48
49impl Solution {
50 fn delete_node(root: TreeLink, key: i32) -> TreeLink {
51 root.delete(key)
52 }
53}
54
55#[test]
56fn test() {
57 let root = tree!(5, tree!(3, tree!(2), tree!(4)), tree!(6, None, tree!(7)));
58 let key = 3;
59 let res = tree!(5, tree!(4, tree!(2), None), tree!(6, None, tree!(7)));
60 assert_eq!(Solution::delete_node(root, key), res);
61}