https://atcoder.jp/contests/abc308/tasks/abc308_e
Aの組合せで分けて考えると楽になります。
// MEX #![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() } //////////////////// process //////////////////// fn read_input() -> (Vec<i32>, String) { let _N: usize = read(); let A = read_vec(); let S = read(); (A, S) } fn mex(triplet: (i32, i32, i32)) -> i32 { for n in 0..3 { if triplet.iter().all(|&m| n != m) { return n } } 3 } fn g(t: (i32, i32, i32), A: &Vec<i32>, S: &String) -> i64 { let N = A.len(); let mut v: Vec<i64> = vec![0; N+1]; for i in (0..N).rev() { if S[ } } fn f(A: Vec<i32>, S: String) -> i64 { let cs: Vec<char> = S.chars().collect(); let triplets: Vec<(i32, i32, i32)> = (0..27).map(|n| (n%3, n/3%3, n/9)); triplets.into_iter().map(|t| (mex(t) as i64) * g(t, &A, &cs)).sum::<i64>() } //////////////////// main //////////////////// fn main() { let (A, S) = read_input(); println!("{}", f(A, S)) }