algorithms_fourth 0.1.10

//!
//! ## 使用
//! ```
//!     use std::fs;
//!
//!     use rand::seq::SliceRandom;
//!
//!     use algorithms_fourth::search::{
//!         binary_search::BinarySearchST,
//!         binary_tree_search::BST,
//!         linear_probing_hash::LinearProbingHashST,
//!         red_black_search::trace::TraceRedBlackBST,
//!         red_black_search::{trace, RedBlackBST},
//!         separate_chaining_hash::SeparateChainingHashST,
//!         sequential::SequentialSearchST,
//!         ST,
//!     };
//!
//!         let mut st: Box<TraceRedBlackBST<usize, usize>> =
//!             Box::new(trace::new(RedBlackBST::new(), vec![]));
//!         let mut rng = rand::thread_rng();
//!         let n = 23;
//!         let mut arr: Vec<usize> = (1..=n).collect();
//!         arr.shuffle(&mut rng);
//!         for (i, key) in arr.iter().enumerate() {
//!             st.put(*key, key + 10);
//!             assert_eq!(st.size(), i + 1);
//!             assert!(st.is_balanced());
//!         }
//!         let mut arr: Vec<usize> = (1..=n).collect();
//!         arr.shuffle(&mut rng);
//!         for (i, key) in arr.iter().enumerate() {
//!             st.delete(key);
//!             // let _ = fs::write("trace.dot", st.get_graph());
//!             assert!(st.is_balanced());
//!             assert_eq!(st.size(), arr.len() - i - 1);
//!         }
//!         //let _ = fs::write("trace.dot", st.get_graph());
//!         println!("{}",st.get_graph());
//! ```
use std::cmp::Ordering;

use super::ST;

pub mod trace;
///
#[derive(Default)]
pub struct RedBlackBST<K: PartialOrd, V> {
    root: Link<K, V>,
}
enum Color {
    Red,
    Black,
}
impl Color {
    fn flip(&self) -> Self {
        match self {
            Color::Red => Color::Black,
            Color::Black => Color::Red,
        }
    }
}
impl Clone for Color {
    fn clone(&self) -> Self {
        match self {
            Self::Red => Self::Red,
            Self::Black => Self::Black,
        }
    }
}
struct Node<K, V> {
    key: K,
    val: V,
    left: Link<K, V>,
    right: Link<K, V>,
    n: usize,
    color: Color,
}
impl<K, V> Node<K, V> {
    fn update_n(&mut self) {
        let mut n = 1;
        if self.left.is_some() {
            n += self.left.as_deref().unwrap().n;
        }
        if self.right.is_some() {
            n += self.right.as_deref().unwrap().n;
        }
        self.n = n;
    }
    fn new(key: K, val: V, n: usize, color: Color) -> Box<Self> {
        Box::new(Node {
            key,
            val,
            left: None,
            right: None,
            n,
            color,
        })
    }
    fn is_left_red(&self) -> bool {
        self.left
            .as_deref()
            .map_or_else(|| false, |left| left.is_red())
    }
    fn is_left_red_of_right(&self) -> bool {
        self.right.as_deref().map_or(false, |f| f.is_left_red())
    }
    fn is_right_red(&self) -> bool {
        self.right
            .as_deref()
            .map_or_else(|| false, |node| node.is_red())
    }
    fn is_red(&self) -> bool {
        match self.color {
            Color::Red => true,
            Color::Black => false,
        }
    }
    fn rotate_left(mut self) -> Box<Node<K, V>> {
        let mut x = self.right.take().unwrap();
        self.right = x.left.take();
        x.left = Some(Box::new(self));
        let left = x.left.as_deref_mut().unwrap();
        x.color = left.color.clone();
        left.color = Color::Red;
        // 只是变换了子树的根节点，所以子树的n不变
        x.n = left.n;
        left.update_n();
        x
    }
    /// 使得left与self组成>2节点
    fn rotate_right(mut self) -> Box<Self> {
        if self.left.is_none() {
            // 右旋的前提是有左节点
            return Box::new(self);
        }
        let mut x = self.left.take().unwrap();
        self.left = x.right.take();
        x.right = Some(Box::new(self));
        let right = x.right.as_deref_mut().unwrap();
        x.color = right.color.clone();
        right.color = Color::Red;
        x.n = right.n;
        right.update_n();
        x
    }
    /// 翻转节点颜色，使得要么当前节点与其左右子节点组成4节点，
    /// 要么当前节点解除与其左右子节点的组合，然后与其父节点组成>2节点
    fn flip_colors(&mut self) {
        self.color = self.color.flip();
        let left = self
            .left
            .as_deref_mut()
            .expect("翻转颜色的前提是子节点非none，left不符合");
        left.color = left.color.flip();
        let right = self
            .right
            .as_deref_mut()
            .expect("翻转颜色的前提是子节点非none，right不符合");
        right.color = right.color.flip();
    }

    fn need_rotate_left(&self) -> bool {
        self.is_right_red() && !self.is_left_red()
    }
    fn need_rotate_right(&self) -> bool {
        self.is_left_red() && self.left.as_deref().map_or(false, |f| f.is_left_red())
    }
    fn need_flip_colors(&self) -> bool {
        self.is_left_red() && self.is_right_red()
    }
    /// 使得left与self组成>2节点
    fn move_red_left(mut self: Box<Self>) -> Box<Self> {
        self.flip_colors();
        if self.is_left_red_of_right() {
            //  此时self是5节点，包含left、self、right、right.left,
            //  目标是left与self组成3节点，并且以right为根节点的子树仍然保持红黑树规则，
            //  所以right右旋，使得4个点是从左到右的关系，
            //  然后变换根为right，翻转根的颜色，达到目的
            self.right = Some(self.right.unwrap().rotate_right());
            //   通过rotate_left使得self变为right
            self = self.rotate_left();
            self.flip_colors();
        }
        self
    }
    fn move_red_right(mut self: Box<Self>) -> Box<Self> {
        self.flip_colors();
        if self.left.as_deref().map_or(false, |f| f.is_left_red()) {
            self = self.rotate_right();
            self.flip_colors();
        }
        self
    }
}
type Link<K, V> = Option<Box<Node<K, V>>>;
struct Temp<K, V> {
    /// 转换后的子树根节点
    node: Link<K, V>,
    /// 从子树中删除的节点
    target: Box<Node<K, V>>,
}
impl<K: PartialOrd, V> RedBlackBST<K, V> {
    pub fn new() -> Self {
        RedBlackBST { root: None }
    }
    fn size_for(node: &Link<K, V>) -> usize {
        match node {
            Some(t) => t.n,
            None => 0,
        }
    }
    pub fn is_balanced(&self) -> bool {
        let mut black = 0;
        let mut x = self.root.as_deref();
        while let Some(t) = x {
            if !t.is_red() {
                black += 1;
            }
            x = t.left.as_deref();
        }
        Self::is_balanced_inner(&self.root, black)
    }
    fn is_balanced_inner(node: &Link<K, V>, mut black: usize) -> bool {
        if let Some(x) = node {
            if !x.is_red() {
                if black == 0 {
                    return false;
                }
                black -= 1;
            }
            Self::is_balanced_inner(&x.left, black) && Self::is_balanced_inner(&x.right, black)
        } else {
            black == 0
        }
    }
    fn get_from<'a>(node: &'a Link<K, V>, key: &K) -> Option<&'a V> {
        match node {
            Some(t) => match key.partial_cmp(&t.key).unwrap() {
                Ordering::Less => Self::get_from(&t.left, key),
                Ordering::Equal => Some(&t.val),
                Ordering::Greater => Self::get_from(&t.right, key),
            },
            None => None,
        }
    }

    fn put_for(node: Link<K, V>, key: K, val: V) -> Link<K, V> {
        let mut target = match node {
            Some(mut t) => match &key.partial_cmp(&t.key).unwrap() {
                Ordering::Less => {
                    t.as_mut().left = RedBlackBST::put_for(t.left.take(), key, val);
                    t.update_n();
                    Some(t)
                }
                Ordering::Equal => {
                    t.as_mut().val = val;
                    Some(t)
                }
                Ordering::Greater => {
                    t.as_mut().right = RedBlackBST::put_for(t.right.take(), key, val);
                    t.update_n();
                    Some(t)
                }
            },
            None => Some(Node::new(key, val, 1, Color::Red)),
        };
        if target.is_some() {
            let node = target.as_deref().unwrap();
            if node.need_rotate_left() {
                target = Some(target.unwrap().rotate_left());
            }
            let node = target.as_deref().unwrap();
            if node.need_rotate_right() {
                target = Some(target.unwrap().rotate_right());
            }
            let node = target.as_deref().unwrap();
            if node.need_flip_colors() {
                target.as_deref_mut().unwrap().flip_colors();
            }
            target.as_deref_mut().unwrap().update_n();
        }
        target
    }
    fn min_for(node: &Link<K, V>) -> Option<&K> {
        node.as_deref().and_then(|t| {
            if t.left.is_some() {
                RedBlackBST::min_for(&t.left)
            } else {
                Some(&t.key)
            }
        })
    }
    fn balance_delete_min(node: Link<K, V>) -> Temp<K, V> {
        let mut h = node.expect("balance delete min have to has none node");
        let left = &h.as_ref().left;
        //  已经是最小节点
        if left.is_none() {
            return Temp {
                node: None,
                target: h,
            };
        }
        if !(h.is_left_red()) && !(h.left.as_deref().map_or(false, |f| f.is_left_red())) {
            h = h.move_red_left();
        }
        let Temp { node, target } = RedBlackBST::balance_delete_min(h.left.take());
        h.left = node;
        h = RedBlackBST::balance(h);
        Temp {
            node: Some(h),
            target,
        }
    }
    fn max_for(node: &Link<K, V>) -> Option<&K> {
        node.as_deref().and_then(|t| {
            if t.right.is_some() {
                RedBlackBST::max_for(&t.right)
            } else {
                Some(&t.key)
            }
        })
    }
    /// ## 条件：
    ///  1. node 一定是3节点的一员(或是树中唯一节点)<br/>
    ///  2. 以node为根的子树仍然是满足红黑树规则<br/>
    /// ## 情况(始终以2-3树来看结构):<br/>
    ///
    ///  a. node为目标节点,删除其右子树的最小节点，并用之替换node的信息<br/>
    ///  b. node大于目标节点:<br/>
    ///     ba. node当前为黑色节点，则其左节点一定为红节点，满足条件1，可以移到子树<br/>
    ///     bb. node当前为红色节点，则其左节点一定黑色节点:<br/>
    ///         bba. 如果其左节点的左节点是红节点，则满足条件1，可以移到子树<br/>
    ///         bbb. 如果其左节点的左节点是黑节点，则需要通过变换使得左节点变为红色节点:<br/>
    ///             bbba. 如果其右节点的左节点是红色节点<br/>
    ///             bbbb. 否<br/>
    ///  c. node小于目标节点(条件2保证了其右节点一定是黑色节点):<br/>
    ///     ca. 其右节点的左节点为红色,满足条件1<br/>
    ///     cb. 否则，如果node非红，则需要先将node变为红色节点，然后将其右节点变为红色节点，满足条件1<br/>
    fn delete_for(node: Link<K, V>, key: &K) -> Link<K, V> {
        let mut t = node.expect("如果key在子树中，那么子树一定非none");
        let ans = match key.partial_cmp(&t.key).unwrap() {
            Ordering::Less => {
                let left = t
                    .left
                    .as_deref()
                    .expect("因为小于根节点，所以目标节点一定在left，所以left一定非none");
                if !left.is_red() && !left.is_left_red() {
                    // 此时下一个要搜索的节点为2节点，需要将其压为3/4节点
                    t = t.move_red_left();
                }
                t.left = RedBlackBST::delete_for(t.left, key);
                Some(t)
            }
            _ => {
                // 根据红黑树规则，只可能left为red
                if t.is_left_red() {
                    // 因为用的删除规则是从右子树中选取最小节点来代替被删节点，
                    // 所以这里要保证在节点被删后不会出现空节点，必须满足以下两个条件之一：
                    // 1. 有右节点
                    // 2. 同时没有左右节点
                    // 在没有右节点的情况下，这里通过右旋使得条件2得到满足。
                    // 在有右节点的情况下，
                    //   1. 如果当前节点是目标节点，则右子树需要做删除最小节点操作，为了平衡性，当前节点做颜色翻转动作，而翻转的前提是当前节点是红节点，所以需要右旋。
                    //   2. 如果当前节点不是目标节点，则需要走到右节点。为了保证右节点有能力是>2节点的一员，还是需要右旋，使得当前节点变成>2节点的右键，以便之后有翻转能力
                    t = t.rotate_right();
                    // 此时以t为根节点的子树依然满足红黑树规则
                }
                if key.partial_cmp(&t.key).unwrap().is_eq() && t.right.is_none() {
                    // 根据红黑树规则，此时left一定是none，该节点可以直接删除
                    return None;
                }
                let right = t
                    .right
                    .as_ref()
                    .expect("此时目标节点一定在right子树中，所以right不为none");
                // 当右旋后right可能为红色
                if !right.is_red() && !right.is_left_red() {
                    // 次数确认可疑的目标节点right一定是一个2节点，所以需要做变换来保证right在>2节点中
                    t = t.move_red_right();
                }
                // 之前已经排除了目标节点在叶子结点中的情况,也保证了子树中删除最小节点的条件
                // ：该子树的根节点为红色，有翻转的能力
                if key.partial_cmp(&t.key).unwrap().is_eq() {
                    let Temp { node, mut target } = RedBlackBST::balance_delete_min(t.right.take());
                    target.color = t.color;
                    target.left = t.left.take();
                    target.right = node;
                    t = target
                } else {
                    t.right = RedBlackBST::delete_for(t.right.take(), key);
                }
                Some(t)
            }
        }
        .map(|f| RedBlackBST::balance(f));
        ans
    }
    fn balance(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
        if node.need_rotate_left() {
            node = node.rotate_left();
        }
        if node.need_rotate_right() {
            node = node.rotate_right();
        }
        if node.need_flip_colors() {
            node.flip_colors();
        }
        node.update_n();
        node
    }
    fn rank_for(node: &Link<K, V>, key: &K) -> usize {
        if let Some(t) = node {
            match key.partial_cmp(&t.key).unwrap() {
                Ordering::Less => RedBlackBST::rank_for(&t.left, key),
                Ordering::Equal => {
                    if t.left.is_some() {
                        t.left.as_deref().unwrap().n
                    } else {
                        0
                    }
                }
                Ordering::Greater => {
                    if t.right.is_some() {
                        let mut n = 1 + RedBlackBST::rank_for(&t.right, key);
                        if t.left.is_some() {
                            n += t.left.as_deref().unwrap().n;
                        }
                        n
                    } else {
                        t.n
                    }
                }
            }
        } else {
            0
        }
    }
    fn select_for(node: &Link<K, V>, k: usize) -> Option<&K> {
        node.as_deref().and_then(|t| {
            let mut n = 0;
            if let Some(left) = &t.left {
                n += left.n;
            }
            match k.partial_cmp(&n).unwrap() {
                Ordering::Less => RedBlackBST::select_for(&t.left, k),
                Ordering::Equal => Some(&t.key),
                Ordering::Greater => RedBlackBST::select_for(&t.right, k - n - 1),
            }
        })
    }
    fn floor_for<'a>(node: &'a Link<K, V>, key: &K) -> Option<&'a K> {
        node.as_deref().and_then(|t| -> Option<&K> {
            match &t.key.partial_cmp(key).unwrap() {
                Ordering::Less => RedBlackBST::floor_for(&t.right, key).or(Some(&t.key)),
                Ordering::Equal => Some(&t.key),
                Ordering::Greater => RedBlackBST::floor_for(&t.left, key),
            }
        })
    }
    fn ceiling_for<'a>(node: &'a Link<K, V>, key: &K) -> Option<&'a K> {
        node.as_deref().and_then(|t| -> Option<&K> {
            match &t.key.partial_cmp(key).unwrap() {
                Ordering::Less => RedBlackBST::ceiling_for(&t.right, key),
                Ordering::Equal => Some(&t.key),
                Ordering::Greater => RedBlackBST::ceiling_for(&t.left, key).or(Some(&t.key)),
            }
        })
    }
}
impl<K: PartialOrd, V> ST<K, V> for RedBlackBST<K, V> {
    fn put(&mut self, key: K, value: V) {
        self.root = RedBlackBST::put_for(self.root.take(), key, value);
        if let Some(it) = self.root.as_mut() {
            it.color = Color::Black;
        }
    }

    fn get(&self, key: &K) -> Option<&V> {
        RedBlackBST::get_from(&self.root, key)
    }

    fn delete(&mut self, key: &K) {
        // 因为删除涉及平衡性，代价大，所以先过滤不必要的操作
        if !self.contains(key) {
            return;
        }
        let root = self.root.as_deref_mut().unwrap();
        // 如果根节点为二节点，则通过变色使得根节点后边可以压为三、四节点
        if !root.is_left_red() && !root.is_right_red() {
            root.color = Color::Red;
        }
        self.root = RedBlackBST::delete_for(self.root.take(), key);
        // 与压缩相对应的展开
        if let Some(root) = self.root.as_deref_mut() {
            root.color = Color::Black;
        }
    }

    fn is_empty(&self) -> bool {
        self.root.is_none()
    }

    fn size(&self) -> usize {
        RedBlackBST::size_for(&self.root)
    }

    fn min(&self) -> Option<&K> {
        RedBlackBST::min_for(&self.root)
    }

    fn max(&self) -> Option<&K> {
        RedBlackBST::max_for(&self.root)
    }

    fn floor(&self, key: &K) -> Option<&K> {
        RedBlackBST::floor_for(&self.root, key)
    }

    fn ceiling(&self, key: &K) -> Option<&K> {
        RedBlackBST::ceiling_for(&self.root, key)
    }

    fn rank(&self, key: &K) -> usize {
        RedBlackBST::rank_for(&self.root, key)
    }

    fn select(&self, k: usize) -> Option<&K> {
        RedBlackBST::select_for(&self.root, k)
    }
}