algorithms_fourth 0.1.10

use std::cmp::Ordering;

use super::ST;

#[derive(Default)]
pub struct BST<K: PartialOrd, V> {
    root: Link<K, V>,
}
struct Node<K, V> {
    key: K,
    val: V,
    left: Link<K, V>,
    right: Link<K, V>,
    n: usize,
}
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;
    }
}
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> BST<K, V> {
    fn size_for(node: &Link<K, V>) -> usize {
        match node {
            Some(t) => t.n,
            None => 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> {
        match node {
            Some(mut t) => match &key.partial_cmp(&t.key).unwrap() {
                Ordering::Less => {
                    t.as_mut().left = BST::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 = BST::put_for(t.right.take(), key, val);
                    t.update_n();
                    Some(t)
                }
            },
            None => Some(Box::new(Node {
                key,
                val,
                left: None,
                right: None,
                n: 1,
            })),
        }
    }
    fn min_for(node: &Link<K, V>) -> Option<&K> {
        node.as_deref().and_then(|t| {
            if t.left.is_some() {
                BST::min_for(&t.left)
            } else {
                Some(&t.key)
            }
        })
    }
    fn max_for(node: &Link<K, V>) -> Option<&K> {
        node.as_deref().and_then(|t| {
            if t.right.is_some() {
                BST::max_for(&t.right)
            } else {
                Some(&t.key)
            }
        })
    }
    fn max_node_for(node: Link<K, V>) -> Option<Temp<K, V>> {
        node.and_then(|mut t| {
            if t.right.is_some() {
                if let Some(Temp { node, target }) = BST::max_node_for(t.right) {
                    t.right = node;
                    t.update_n();
                    Some(Temp {
                        node: Some(t),
                        target,
                    })
                } else {
                    None
                }
            } else {
                Some(Temp {
                    node: t.left.take(),
                    target: t,
                })
            }
        })
    }

    fn delete_for(node: Link<K, V>, key: &K) -> Link<K, V> {
        if let Some(mut t) = node {
            match key.partial_cmp(&t.key).unwrap() {
                Ordering::Less => {
                    // println!("less");
                    t.left = BST::delete_for(t.left, key);
                    t.update_n();
                    Some(t)
                }
                Ordering::Equal => {
                    // println!("equal");
                    let right = t.right.take();
                    if let Some(Temp { node, target }) = BST::max_node_for(t.left.take()) {
                        let mut new_node = target;
                        new_node.left = node;
                        new_node.right = right;
                        new_node.update_n();
                        Some(new_node)
                    } else {
                        right
                    }
                }
                Ordering::Greater => {
                    // println!("greater");
                    t.right = BST::delete_for(t.right.take(), key);
                    t.update_n();
                    Some(t)
                }
            }
        } else {
            None
        }
    }
    fn rank_for(node: &Link<K, V>, key: &K) -> usize {
        if let Some(t) = node {
            match key.partial_cmp(&t.key).unwrap() {
                Ordering::Less => BST::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 + BST::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 => BST::select_for(&t.left, k),
                Ordering::Equal => Some(&t.key),
                Ordering::Greater => BST::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 => BST::floor_for(&t.right, key).or(Some(&t.key)),
                Ordering::Equal => Some(&t.key),
                Ordering::Greater => BST::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 => BST::ceiling_for(&t.right, key),
                Ordering::Equal => Some(&t.key),
                Ordering::Greater => BST::ceiling_for(&t.left, key).or(Some(&t.key)),
            }
        })
    }
}
impl<K: PartialOrd, V> ST<K, V> for BST<K, V> {
    fn put(&mut self, key: K, value: V) {
        self.root = BST::put_for(self.root.take(), key, value)
    }

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

    fn delete(&mut self, key: &K) {
        self.root = BST::delete_for(self.root.take(), key)
    }

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

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

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

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

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

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

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

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