使用rust实现红黑树

下面是一个简单的红黑树实现，使用了Rust语言。

use std::cmp::Ordering;

enum Color {
    Red,
    Black,
}

struct Node<K, V> {
    key: K,
    value: V,
    color: Color,
    left: Option<Box<Node<K, V>>>,
    right: Option<Box<Node<K, V>>>,
}

impl<K: Ord, V> Node<K, V> {
    fn new(key: K, value: V, color: Color) -> Self {
        Node {
            key,
            value,
            color,
            left: None,
            right: None,
        }
    }

    fn is_red(&self) -> bool {
        self.color == Color::Red
    }

    fn rotate_left(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
        let mut right = node.right.take().unwrap();
        node.right = right.left.take();
        right.left = Some(node);
        right.color = right.left.as_ref().unwrap().color;
        right.left.as_mut().unwrap().color = Color::Red;
        right
    }

    fn rotate_right(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
        let mut left = node.left.take().unwrap();
        node.left = left.right.take();
        left.right = Some(node);
        left.color = left.right.as_ref().unwrap().color;
        left.right.as_mut().unwrap().color = Color::Red;
        left
    }

    fn flip_colors(node: &mut Box<Node<K, V>>) {
        node.color = Color::Red;
        node.left.as_mut().unwrap().color = Color::Black;
        node.right.as_mut().unwrap().color = Color::Black;
    }

    fn insert(mut node: Box<Node<K, V>>, key: K, value: V) -> Box<Node<K, V>> {
        match key.cmp(&node.key) {
            Ordering::Less => {
                if node.left.is_none() {
                    node.left = Some(Box::new(Node::new(key, value, Color::Red)));
                } else {
                    node.left = Some(Node::insert(node.left.unwrap(), key, value));
                }
            }
            Ordering::Greater => {
                if node.right.is_none() {
                    node.right = Some(Box::new(Node::new(key, value, Color::Red)));
                } else {
                    node.right = Some(Node::insert(node.right.unwrap(), key, value));
                }
            }
            Ordering::Equal => {
                node.value = value;
            }
        }

        if node.right.as_ref().map_or(false, |x| x.is_red())
            && !node.left.as_ref().map_or(false, |x| x.is_red())
        {
            node = Node::rotate_left(node);
        }

        if node.left.as_ref().map_or(false, |x| x.left.as_ref().map_or(false, |y| y.is_red()))
            && node.right.as_ref().map_or(false, |x| x.is_red())
        {
            node = Node::rotate_right(node);
        }

        if node.left.as_ref().map_or(false, |x| x.is_red())
            && node.right.as_ref().map_or(false, |x| x.is_red())
        {
            Node::flip_colors(&mut node);
        }

        node
    }
}

pub struct RedBlackTree<K: Ord, V> {
    root: Option<Box<Node<K, V>>>,
}

impl<K: Ord, V> RedBlackTree<K, V> {
    pub fn new() -> Self {
        RedBlackTree { root: None }
    }

    pub fn insert(&mut self, key: K, value: V) {
        if self.root.is_none() {
            self.root = Some(Box::new(Node::new(key, value, Color::Black)));
        } else {
            self.root = Some(Node::insert(self.root.unwrap(), key, value));
            self.root.as_mut().unwrap().color = Color::Black;
        }
    }

    pub fn get(&self, key: &K) -> Option<&V> {
        let mut current = self.root.as_ref();
        while let Some(node) = current {
            match key.cmp(&node.key) {
                Ordering::Less => current = node.left.as_ref(),
                Ordering::Greater => current = node.right.as_ref(),
                Ordering::Equal => return Some(&node.value),
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert() {
        let mut tree = RedBlackTree::new();
        tree.insert(3, "three");
        tree.insert(2, "two");
        tree.insert(1, "one");
        tree.insert(4, "four");
        assert_eq!(tree.get(&1), Some(&"one"));
        assert_eq!(tree.get(&2), Some(&"two"));
        assert_eq!(tree.get(&3), Some(&"three"));
        assert_eq!(tree.get(&4), Some(&"four"));
    }
}

这个实现使用了一个枚举类型表示颜色，包含两个值：红和黑。Node结构体表示红黑树的一个节点，包含键、值、颜色、左子树和右子树。Node实现了一些方法，包括is_red()、rotate_left()、rotate_right()和flip_colors()，这些方法用于调整红黑树的平衡。RedBlackTree结构体表示整个红黑树，包含根节点。RedBlackTree实现了一些方法，包括new()、insert()和get()，这些方法用于创建、插入和获取红黑树的节点。