first-rust/src/bst.rs
2024-11-22 17:41:54 -06:00

109 lines
2.6 KiB
Rust

use std::cmp::Ordering::*;
#[derive(Debug)]
pub struct BinarySearchTree<T: Ord> {
root: Option<Box<Node<T>>>,
}
#[derive(Debug)]
struct Node<T> {
data: T,
left: Option<Box<Node<T>>>,
right: Option<Box<Node<T>>>,
}
impl<T: Ord> BinarySearchTree<T> {
pub fn new() -> BinarySearchTree<T> {
BinarySearchTree { root: None }
}
pub fn insert(&mut self, data: T) {
match self.root.take() {
None => {
self.root = Some(Box::new(Node {
data,
left: None,
right: None,
}))
}
Some(mut root) => {
root.insert_child(data);
self.root = Some(root);
}
}
}
pub fn find(&self, data: T) -> Option<T> {
self.root.as_ref().map(|root| root.find(data))?
}
}
impl<T: Ord> Node<T> {
fn new(data: T) -> Self {
Node {
data,
left: None,
right: None,
}
}
fn insert_child(&mut self, data: T) {
match self.data.cmp(&data) {
Less => match self.right.take() {
None => self.right = Some(Box::new(Node::new(data))),
Some(mut right) => {
right.insert_child(data);
self.right = Some(right);
}
},
Equal | Greater => match self.left.take() {
None => self.left = Some(Box::new(Node::new(data))),
Some(mut left) => {
left.insert_child(data);
self.left = Some(left);
}
},
}
}
fn find(&self, data: T) -> Option<T> {
match self.data.cmp(&data) {
Less => self.right.as_ref().map(|right| right.find(data))?,
Equal => Some(data),
Greater => self.left.as_ref().map(|left| left.find(data))?,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn insert_and_find() {
let mut tree = BinarySearchTree::new();
tree.insert(1);
assert_eq!(tree.find(1), Some(1));
}
#[test]
fn insert_three_values_find_third() {
let mut tree = BinarySearchTree::new();
for num in [1, 2, 3] {
tree.insert(num);
}
assert_eq!(tree.find(3), Some(3));
}
#[test]
fn multiple_finds() {
let mut tree = BinarySearchTree::new();
for num in [1, 2, 3] {
tree.insert(num);
}
for num in [1, 2, 3] {
assert_eq!(tree.find(num), Some(num));
}
}
}