AtCoder Beginner Contest 353 E

https://atcoder.jp/contests/abc353/tasks/abc353_e

木を作ればよいです。例1なら、

a(3)┓
    ┣b(2)┓
    ┃    ┗c(1)
    ┗rc(1)

のような木を作ります。数字は文字列の個数です。
Pythonで書いてAIにRustにしてもらいましたが、間違いが多くて修正するのが大変でした。

// Yet Another Sigma Problem
#![allow(non_snake_case)]

use std::fmt;
use std::cmp::min;


//////////////////// 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()
}


//////////////////// Node ////////////////////

struct Node {
    str: String,
    num: i64,
    children: Vec<Node>,
}

impl Node {
    fn new(s: &str) -> Node {
        Node {
            str: s.to_string(),
            num: 1,
            children: Vec::new(),
        }
    }
    
    fn insert(&mut self, s: &str) {
        self.num += 1;
        let mut length = std::usize::MAX;
        let common_length = min(s.len(), self.str.len());
        for i in 0..common_length {
            if s.as_bytes()[i] != self.str.as_bytes()[i] {
                length = i;
                break
            }
        }
        
        if length == std::usize::MAX {
            if s.len() == self.str.len() {
                return
            }
            else if s.len() < self.str.len() {
                let old_str = self.str.clone();
                self.str = s.to_string();
                let mut node = Node::new(&old_str[s.len()..]);
                node.num = self.num - 1;
                node.children.append(&mut self.children);
                self.children = vec![node];
            }
            else {
                let subs = &s[self.str.len()..];
                for child in &mut self.children {
                    if child.str.as_bytes()[0] == subs.as_bytes()[0] {
                        child.insert(subs);
                        return;
                    }
                }
                self.children.push(Node::new(subs));
            }
        }
        else {
            let old_str = self.str.clone();
            self.str = s[..length].to_string();
            let mut node1 = Node::new(&old_str[length..]);
            node1.num = self.num - 1;
            node1.children.append(&mut self.children);
            let node2 = Node::new(&s[length..]);
            self.children = vec![node1, node2];
        }
    }
    
    fn sum(&self) -> i64 {
        let s1 = self.str.len() as i64 * self.num * (self.num - 1) / 2;
        let s2: i64 = self.children.iter().map(|c| c.sum()).sum();
        s1 + s2
    }
}

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s = format!("{}({})", self.str, self.num);
        let v: Vec<String> = self.children.iter()
            .flat_map(|c| format!("{}", c).split('\n').map(|s1| format!("  {}", s1)).collect::<Vec<String>>())
            .collect();
        write!(f, "{}\n{}", s, v.join("\n"))
    }
}


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

struct Tree {
    children: Vec<Node>
}

impl Tree {
    fn new() -> Tree {
        Tree { children: vec![] }
    }
    
    fn insert(&mut self, s: String) {
        for c in self.children.iter_mut() {
            if c.str.as_bytes()[0] == s.as_bytes()[0] {
                c.insert(&s);
                return
            }
        }
        self.children.push(Node::new(&s))
    }
    
    fn sum(&self) -> i64 {
        self.children.iter().map(|c| c.sum()).sum::<i64>()
    }
}

impl fmt::Display for Tree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s = "*\n";
        let v: Vec<String> = self.children.iter()
            .flat_map(|c| format!("{}", c).split('\n').map(|s1| format!("  {}", s1)).collect::<Vec<String>>())
            .collect();
        write!(f, "{}{}", s, v.join("\n"))
    }
}


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

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

fn F(A: Vec<String>) -> i64 {
    let mut tree = Tree::new();
    for s in A.into_iter() {
        tree.insert(s);
    }
    tree.sum()
}

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