AtCoder Beginner Contest 392 F

https://atcoder.jp/contests/abc392/tasks/abc392_f

平衡木が書ければ解けるのは明らかなのですが、今までAIでもなかなか書けませんでした。今回も最初は書けなかったのですが、Pythonでふつうの木を書いて、これを平衡木にと指定したらPythonでは書けました。しかし、これをRustにしてと言ってもなかなかちゃんと書いてくれない。最近できた深く考えるモードを指定したら2回で正確に翻訳してくれました。

// Insert
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// Tree ////////////////////

use std::cmp::max;

type Link = Option<Box<Node>>;

struct Node {
    value: i32,
    left: Link,
    right: Link,
    height: i32,
    size: i32, // Number of nodes in the subtree rooted at this node
}

impl Node {
    fn new(value: i32) -> Self {
        Node {
            value,
            left: None,
            right: None,
            height: 1,
            size: 1,
        }
    }
}

struct AVLTree {
    root: Link,
}

impl AVLTree {
    fn new() -> Self {
        AVLTree { root: None }
    }

    fn get_height(node: &Link) -> i32 {
        node.as_ref().map_or(0, |n| n.height)
    }

    fn get_size(node: &Link) -> i32 {
        node.as_ref().map_or(0, |n| n.size)
    }

    fn update_node(node: &mut Box<Node>) {
        node.height = 1 + max(Self::get_height(&node.left), Self::get_height(&node.right));
        node.size = 1 + Self::get_size(&node.left) + Self::get_size(&node.right);
    }

    fn get_balance(node: &Box<Node>) -> i32 {
        Self::get_height(&node.left) - Self::get_height(&node.right)
    }

    fn rotate_right(mut y: Box<Node>) -> Box<Node> {
        let mut x = y.left.take().expect("Left child must exist for right rotation");
        let t2 = x.right.take();

        // Perform rotation
        x.right = Some(y);
        x.right.as_mut().unwrap().left = t2;

        // Update nodes
        Self::update_node(x.right.as_mut().unwrap()); // Update y (now x.right)
        Self::update_node(&mut x);                    // Update x

        x
    }

    fn rotate_left(mut x: Box<Node>) -> Box<Node> {
        let mut y = x.right.take().expect("Right child must exist for left rotation");
        let t2 = y.left.take();

        // Perform rotation
        y.left = Some(x);
        y.left.as_mut().unwrap().right = t2;

        // Update nodes
        Self::update_node(y.left.as_mut().unwrap()); // Update x (now y.left)
        Self::update_node(&mut y);                   // Update y

        y
    }

    fn insert_node(
        node: Link,
        value: i32,
        pos: i32,
        current_pos: i32,
    ) -> Link {
        if let Some(mut node_box) = node {
            let left_size = Self::get_size(&node_box.left);
            if pos <= current_pos + left_size {
                node_box.left = Self::insert_node(node_box.left.take(), value, pos, current_pos);
            } else {
                node_box.right = Self::insert_node(
                    node_box.right.take(),
                    value,
                    pos,
                    current_pos + left_size + 1,
                );
            }

            Self::update_node(&mut node_box);

            let balance = Self::get_balance(&node_box);

            // Left Left Case
            if balance > 1 && Self::get_balance(node_box.left.as_ref().unwrap()) >= 0 {
                return Some(Self::rotate_right(node_box));
            }

            // Left Right Case
            if balance > 1 && Self::get_balance(node_box.left.as_ref().unwrap()) < 0 {
                node_box.left = Some(Self::rotate_left(node_box.left.take().unwrap()));
                return Some(Self::rotate_right(node_box));
            }

            // Right Right Case
            if balance < -1 && Self::get_balance(node_box.right.as_ref().unwrap()) <= 0 {
                return Some(Self::rotate_left(node_box));
            }

            // Right Left Case
            if balance < -1 && Self::get_balance(node_box.right.as_ref().unwrap()) > 0 {
                node_box.right = Some(Self::rotate_right(node_box.right.take().unwrap()));
                return Some(Self::rotate_left(node_box));
            }

            Some(node_box)
        } else {
            Some(Box::new(Node::new(value)))
        }
    }

    fn insert(&mut self, value: i32, pos: i32) {
        self.root = Self::insert_node(self.root.take(), value, pos, 0);
    }

    fn in_order_traverse(node: &Link, result: &mut Vec<i32>) {
        if let Some(ref node_box) = node {
            Self::in_order_traverse(&node_box.left, result);
            result.push(node_box.value);
            Self::in_order_traverse(&node_box.right, result);
        }
    }

    fn list(&self) -> Vec<i32> {
        let mut result = Vec::with_capacity(Self::get_size(&self.root) as usize);
        Self::in_order_traverse(&self.root, &mut result);
        result
    }
}


//////////////////// process ////////////////////

fn read_input() -> Vec<usize> {
    let _N: usize = read();
    let P: Vec<usize> = read_vec();
    P
}

fn F(P: Vec<usize>) -> String {
    let mut tree = AVLTree::new();
    for (v, p) in P.iter().enumerate() {
        tree.insert(v as i32 + 1, *p as i32 - 1);
    }
    let v = tree.list();
    v.into_iter().map(|e| e.to_string()).collect::<Vec<_>>().join(" ")
}

fn main() {
    let P = read_input();
    println!("{}", F(P))
}