Data-Structure/Exercise/Next2/Q2.rs

use std::io::{self, BufRead};
use std::cmp::max;
use std::fmt::Write;

struct Node {
    key: i32,
    h: i32,
    left: Option<Box<Node>>,
    right: Option<Box<Node>>,
}

impl Node {
    fn new(key: i32) -> Self {
        Node { key, h: 1, left: None, right: None }
    }
}

fn get_h(n: &Option<Box<Node>>) -> i32 {
    n.as_ref().map_or(0, |node| node.h)
}

fn update_h(n: &mut Node) {
    n.h = max(get_h(&n.left), get_h(&n.right)) + 1;
}

fn rotate_right(mut y: Box<Node>) -> Box<Node> {
    let mut x = y.left.take().unwrap();
    y.left = x.right.take();
    update_h(&mut y);
    x.right = Some(y);
    update_h(&mut x);
    x
}

fn rotate_left(mut x: Box<Node>) -> Box<Node> {
    let mut y = x.right.take().unwrap();
    x.right = y.left.take();
    update_h(&mut x);
    y.left = Some(x);
    update_h(&mut y);
    y
}

fn balance(mut n: Box<Node>, op_name: &str, val: i32, level: i32, msgs: &mut Vec<String>) -> Box<Node> {
    update_h(&mut n);
    let factor = get_h(&n.left) - get_h(&n.right);

    if factor > 1 {
        let left_node = n.left.as_mut().unwrap();
        if get_h(&left_node.left) >= get_h(&left_node.right) {
            msgs.push(format!("{}: Rebalance subtree rooted at node {} on level {} with right rotation. ", op_name, n.key, level));
            return rotate_right(n);
        } else {
            msgs.push(format!("{}: Rebalance subtree rooted at node {} on level {} with left rotation and right rotation. ", op_name, n.key, level));
            n.left = Some(rotate_left(n.left.take().unwrap()));
            return rotate_right(n);
        }
    } else if factor < -1 {
        let right_node = n.right.as_mut().unwrap();
        if get_h(&right_node.right) >= get_h(&right_node.left) {
            msgs.push(format!("{}: Rebalance subtree rooted at node {} on level {} with left rotation. ", op_name, n.key, level));
            return rotate_left(n);
        } else {
            msgs.push(format!("{}: Rebalance subtree rooted at node {} on level {} with right rotation and left rotation. ", op_name, n.key, level));
            n.right = Some(rotate_right(n.right.take().unwrap()));
            return rotate_left(n);
        }
    }
    n
}

fn insert(n: Option<Box<Node>>, key: i32, depth: i32, msgs: &mut Vec<String>) -> Option<Box<Node>> {
    let mut node = match n {
        None => return Some(Box::new(Node::new(key))),
        Some(b) => b,
    };

    if key < node.key {
        node.left = insert(node.left.take(), key, depth + 1, msgs);
    } else if key > node.key {
        node.right = insert(node.right.take(), key, depth + 1, msgs);
    } else {
        return Some(node);
    }
    Some(balance(node, "Insert", key, depth, msgs))
}

fn extract_min(mut n: Box<Node>) -> (Box<Node>, i32) {
    if n.left.is_none() {
        return (n, 0); // Placeholder
    }
    // This part is tricky in Box AVL, simplified:
    let mut curr = &n;
    while curr.left.is_some() { curr = curr.left.as_ref().unwrap(); }
    let v = curr.key;
    (n, v) // Logical placeholder for actual rebalancing removal
}

// 简化逻辑：AVL删除的递归平衡
fn remove(n: Option<Box<Node>>, key: i32, depth: i32, msgs: &mut Vec<String>) -> Option<Box<Node>> {
    let mut node = n?;
    if key < node.key {
        node.left = remove(node.left.take(), key, depth + 1, msgs);
    } else if key > node.key {
        node.right = remove(node.right.take(), key, depth + 1, msgs);
    } else {
        if node.left.is_none() { return node.right; }
        if node.right.is_none() { return node.left; }
        
        let mut min_node = node.right.as_ref().unwrap();
        while let Some(ref l) = min_node.left { min_node = l; }
        let min_val = min_node.key;
        node.key = min_val;
        node.right = remove(node.right.take(), min_val, depth + 1, msgs);
    }
    Some(balance(node, "Remove", key, depth, msgs))
}

fn walk_in(n: &Option<Box<Node>>, res: &mut String) {
    if let Some(ref node) = n {
        walk_in(&node.left, res);
        write!(res, "{} ", node.key).ok();
        walk_in(&node.right, res);
    }
}

fn walk_pre(n: &Option<Box<Node>>, res: &mut String) {
    if let Some(ref node) = n {
        write!(res, "{} ", node.key).ok();
        walk_pre(&node.left, res);
        walk_pre(&node.right, res);
    }
}

fn main() {
    let input = io::stdin();
    let mut lines = input.lock().lines();
    let mut case_idx = 1;

    while let Some(Ok(first_line)) = lines.next() {
        let t: usize = match first_line.trim().parse() {
            Ok(v) => v,
            Err(_) => break,
        };

        let mut root: Option<Box<Node>> = None;
        let mut total_msgs = Vec::new();

        for _ in 0..t {
            let line = lines.next().unwrap().unwrap();
            let pts: Vec<&str> = line.split_whitespace().collect();
            let k: i32 = pts[1].parse().unwrap();
            let mut step_msgs = Vec::new();
            if pts[0] == "Insert" {
                root = insert(root, k, 0, &mut step_msgs);
            } else {
                root = remove(root, k, 0, &mut step_msgs);
            }
            if !step_msgs.is_empty() {
                total_msgs.push(step_msgs.join(""));
            }
        }

        println!("Case {}:", case_idx);
        if !total_msgs.is_empty() {
            for m in total_msgs { println!("{}", m); }
            println!();
        }

        if let Some(_) = root {
            let mut s1 = String::new();
            let mut s2 = String::new();
            walk_in(&root, &mut s1);
            walk_pre(&root, &mut s2);
            println!("{}", s1.trim_end());
            println!("\n{}", s2.trim_end());
        }
        println!();
        case_idx += 1;
    }
}