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