Skip to main content

rustgym/leetcode/
_450_delete_node_in_a_bst.rs

1struct 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}