二叉树

  1. use std::cmp::Ordering;
  2. use std::ops::Deref;
  4. /// This struct implements as Binary Search Tree (BST), which is a
  5. /// simple data structure for storing sorted data
  6. pub struct BinarySearchTree<T>
  7. where
  8.     T: Ord,
  9. {
  10.     value: Option<T>,
  11.     left: Option<Box<BinarySearchTree<T>>>,
  12.     right: Option<Box<BinarySearchTree<T>>>,
  13. }
  15. impl<T> Default for BinarySearchTree<T>
  16. where
  17.     T: Ord,
  18. {
  19.     fn default() -> Self {
  20.         Self::new()
  21.     }
  22. }
  24. impl<T> BinarySearchTree<T>
  25. where
  26.     T: Ord,
  27. {
  28.     /// Create a new, empty BST
  29.     pub fn new() -> BinarySearchTree<T> {
  30.         BinarySearchTree {
  31.             value: None,
  32.             left: None,
  33.             right: None,
  34.         }
  35.     }
  37.     /// Find a value in this tree. Returns True iff value is in this
  38.     /// tree, and false otherwise
  39.     pub fn search(&self, value: &T) -> bool {
  40.         match &self.value {
  41.             Some(key) => {
  42.                 match key.cmp(value) {
  43.                     Ordering::Equal => {
  44.                         // key == value
  45.                         true
  46.                     }
  47.                     Ordering::Greater => {
  48.                         // key > value
  49.                         match &self.left {
  50.                             Some(node) => node.search(value),
  51.                             None => false,
  52.                         }
  53.                     }
  54.                     Ordering::Less => {
  55.                         // key < value
  56.                         match &self.right {
  57.                             Some(node) => node.search(value),
  58.                             None => false,
  59.                         }
  60.                     }
  61.                 }
  62.             }
  63.             None => false,
  64.         }
  65.     }
  67.     /// Returns a new iterator which iterates over this tree in order
  68.     pub fn iter(&self) -> impl Iterator<Item = &T> {
  69.         BinarySearchTreeIter::new(self)
  70.     }
  72.     /// Insert a value into the appropriate location in this tree.
  73.     pub fn insert(&mut self, value: T) {
  74.         if self.value.is_none() {
  75.             self.value = Some(value);
  76.         } else {
  77.             match &self.value {
  78.                 None => (),
  79.                 Some(key) => {
  80.                     let target_node = if value < *key {
  81.                         &mut self.left
  82.                     } else {
  83.                         &mut self.right
  84.                     };
  85.                     match target_node {
  86.                         Some(ref mut node) => {
  87.                             node.insert(value);
  88.                         }
  89.                         None => {
  90.                             let mut node = BinarySearchTree::new();
  91.                             node.insert(value);
  92.                             *target_node = Some(Box::new(node));
  93.                         }
  94.                     }
  95.                 }
  96.             }
  97.         }
  98.     }
  100.     /// Returns the smallest value in this tree
  101.     pub fn minimum(&self) -> Option<&T> {
  102.         match &self.left {
  103.             Some(node) => node.minimum(),
  104.             None => match &self.value {
  105.                 Some(value) => Some(value),
  106.                 None => None,
  107.             },
  108.         }
  109.     }
  111.     /// Returns the largest value in this tree
  112.     pub fn maximum(&self) -> Option<&T> {
  113.         match &self.right {
  114.             Some(node) => node.maximum(),
  115.             None => match &self.value {
  116.                 Some(value) => Some(value),
  117.                 None => None,
  118.             },
  119.         }
  120.     }
  122.     /// Returns the largest value in this tree smaller than value
  123.     pub fn floor(&self, value: &T) -> Option<&T> {
  124.         match &self.value {
  125.             Some(key) => {
  126.                 match key.cmp(value) {
  127.                     Ordering::Greater => {
  128.                         // key > value
  129.                         match &self.left {
  130.                             Some(node) => node.floor(value),
  131.                             None => None,
  132.                         }
  133.                     }
  134.                     Ordering::Less => {
  135.                         // key < value
  136.                         match &self.right {
  137.                             Some(node) => {
  138.                                 let val = node.floor(value);
  139.                                 match val {
  140.                                     Some(_) => val,
  141.                                     None => Some(key),
  142.                                 }
  143.                             }
  144.                             None => Some(key),
  145.                         }
  146.                     }
  147.                     Ordering::Equal => Some(key),
  148.                 }
  149.             }
  150.             None => None,
  151.         }
  152.     }
  154.     /// Returns the smallest value in this tree larger than value
  155.     pub fn ceil(&self, value: &T) -> Option<&T> {
  156.         match &self.value {
  157.             Some(key) => {
  158.                 match key.cmp(value) {
  159.                     Ordering::Less => {
  160.                         // key < value
  161.                         match &self.right {
  162.                             Some(node) => node.ceil(value),
  163.                             None => None,
  164.                         }
  165.                     }
  166.                     Ordering::Greater => {
  167.                         // key > value
  168.                         match &self.left {
  169.                             Some(node) => {
  170.                                 let val = node.ceil(value);
  171.                                 match val {
  172.                                     Some(_) => val,
  173.                                     None => Some(key),
  174.                                 }
  175.                             }
  176.                             None => Some(key),
  177.                         }
  178.                     }
  179.                     Ordering::Equal => {
  180.                         // key == value
  181.                         Some(key)
  182.                     }
  183.                 }
  184.             }
  185.             None => None,
  186.         }
  187.     }
  188. }
  190. struct BinarySearchTreeIter<'a, T>
  191. where
  192.     T: Ord,
  193. {
  194.     stack: Vec<&'a BinarySearchTree<T>>,
  195. }
  197. impl<'a, T> BinarySearchTreeIter<'a, T>
  198. where
  199.     T: Ord,
  200. {
  201.     pub fn new(tree: &BinarySearchTree<T>) -> BinarySearchTreeIter<T> {
  202.         let mut iter = BinarySearchTreeIter { stack: vec![tree] };
  203.         iter.stack_push_left();
  204.         iter
  205.     }
  207.     fn stack_push_left(&mut self) {
  208.         while let Some(child) = &self.stack.last().unwrap().left {
  209.             self.stack.push(child);
  210.         }
  211.     }
  212. }
  214. impl<'a, T> Iterator for BinarySearchTreeIter<'a, T>
  215. where
  216.     T: Ord,
  217. {
  218.     type Item = &'a T;
  220.     fn next(&mut self) -> Option<&'a T> {
  221.         if self.stack.is_empty() {
  222.             None
  223.         } else {
  224.             let node = self.stack.pop().unwrap();
  225.             if node.right.is_some() {
  226.                 self.stack.push(node.right.as_ref().unwrap().deref());
  227.                 self.stack_push_left();
  228.             }
  229.             node.value.as_ref()
  230.         }
  231.     }
  232. }
  234. #[cfg(test)]
  235. mod test {
  236.     use super::BinarySearchTree;
  238.     fn prequel_memes_tree() -> BinarySearchTree<&'static str> {
  239.         let mut tree = BinarySearchTree::new();
  240.         tree.insert("hello there");
  241.         tree.insert("general kenobi");
  242.         tree.insert("you are a bold one");
  243.         tree.insert("kill him");
  244.         tree.insert("back away...I will deal with this jedi slime myself");
  245.         tree.insert("your move");
  246.         tree.insert("you fool");
  247.         tree
  248.     }
  250.     #[test]
  251.     fn test_search() {
  252.         let tree = prequel_memes_tree();
  253.         assert!(tree.search(&"hello there"));
  254.         assert!(tree.search(&"you are a bold one"));
  255.         assert!(tree.search(&"general kenobi"));
  256.         assert!(tree.search(&"you fool"));
  257.         assert!(tree.search(&"kill him"));
  258.         assert!(
  259.             !tree.search(&"but i was going to tosche station to pick up some power converters",)
  260.         );
  261.         assert!(!tree.search(&"only a sith deals in absolutes"));
  262.         assert!(!tree.search(&"you underestimate my power"));
  263.     }
  265.     #[test]
  266.     fn test_maximum_and_minimum() {
  267.         let tree = prequel_memes_tree();
  268.         assert_eq!(*tree.maximum().unwrap(), "your move");
  269.         assert_eq!(
  270.             *tree.minimum().unwrap(),
  271.             "back away...I will deal with this jedi slime myself"
  272.         );
  273.         let mut tree2: BinarySearchTree<i32> = BinarySearchTree::new();
  274.         assert!(tree2.maximum().is_none());
  275.         assert!(tree2.minimum().is_none());
  276.         tree2.insert(0);
  277.         assert_eq!(*tree2.minimum().unwrap(), 0);
  278.         assert_eq!(*tree2.maximum().unwrap(), 0);
  279.         tree2.insert(-5);
  280.         assert_eq!(*tree2.minimum().unwrap(), -5);
  281.         assert_eq!(*tree2.maximum().unwrap(), 0);
  282.         tree2.insert(5);
  283.         assert_eq!(*tree2.minimum().unwrap(), -5);
  284.         assert_eq!(*tree2.maximum().unwrap(), 5);
  285.     }
  287.     #[test]
  288.     fn test_floor_and_ceil() {
  289.         let tree = prequel_memes_tree();
  290.         assert_eq!(*tree.floor(&"hello there").unwrap(), "hello there");
  291.         assert_eq!(
  292.             *tree
  293.                 .floor(&"these are not the droids you're looking for")
  294.                 .unwrap(),
  295.             "kill him"
  296.         );
  297.         assert!(tree.floor(&"another death star").is_none());
  298.         assert_eq!(*tree.floor(&"you fool").unwrap(), "you fool");
  299.         assert_eq!(
  300.             *tree.floor(&"but i was going to tasche station").unwrap(),
  301.             "back away...I will deal with this jedi slime myself"
  302.         );
  303.         assert_eq!(
  304.             *tree.floor(&"you underestimate my power").unwrap(),
  305.             "you fool"
  306.         );
  307.         assert_eq!(*tree.floor(&"your new empire").unwrap(), "your move");
  308.         assert_eq!(*tree.ceil(&"hello there").unwrap(), "hello there");
  309.         assert_eq!(
  310.             *tree
  311.                 .ceil(&"these are not the droids you're looking for")
  312.                 .unwrap(),
  313.             "you are a bold one"
  314.         );
  315.         assert_eq!(
  316.             *tree.ceil(&"another death star").unwrap(),
  317.             "back away...I will deal with this jedi slime myself"
  318.         );
  319.         assert_eq!(*tree.ceil(&"you fool").unwrap(), "you fool");
  320.         assert_eq!(
  321.             *tree.ceil(&"but i was going to tasche station").unwrap(),
  322.             "general kenobi"
  323.         );
  324.         assert_eq!(
  325.             *tree.ceil(&"you underestimate my power").unwrap(),
  326.             "your move"
  327.         );
  328.         assert!(tree.ceil(&"your new empire").is_none());
  329.     }
  331.     #[test]
  332.     fn test_iterator() {
  333.         let tree = prequel_memes_tree();
  334.         let mut iter = tree.iter();
  335.         assert_eq!(
  336.             iter.next().unwrap(),
  337.             &"back away...I will deal with this jedi slime myself"
  338.         );
  339.         assert_eq!(iter.next().unwrap(), &"general kenobi");
  340.         assert_eq!(iter.next().unwrap(), &"hello there");
  341.         assert_eq!(iter.next().unwrap(), &"kill him");
  342.         assert_eq!(iter.next().unwrap(), &"you are a bold one");
  343.         assert_eq!(iter.next().unwrap(), &"you fool");
  344.         assert_eq!(iter.next().unwrap(), &"your move");
  345.         assert_eq!(iter.next(), None);
  346.         assert_eq!(iter.next(), None);
  347.     }
  348. }