AtCoder Beginner Contest 421 E

https://atcoder.jp/contests/abc421/tasks/abc421_e

今のダイスの状態をキープしたダイスは面の数を負にした数で表します。また、ソートして状態数を減らします。マスクが32通りあるので、その中で最大値を選びます。逆から辿って解答を得ます。

// Yacht
#![allow(non_snake_case)]

use std::collections::{HashSet, HashMap};


//////////////////// 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 ////////////////////

type State = (i32, i32, i32, i32, i32);
type Probs = HashMap<State, f64>;

fn state_to_array(s: &State) -> [i32; 5] {
    [s.0, s.1, s.2, s.3, s.4]
}

fn array_to_state(a: &[i32; 5]) -> State {
    (a[0], a[1], a[2], a[3], a[4])
}

fn product(n: usize, A: &Vec<i32>) -> Vec<Vec<i32>> {
    if n == 1 {
        A.iter().map(|&a| vec![a]).collect::<Vec<Vec<i32>>>()
    }
    else if n % 2 == 1 {
        let mut v: Vec<Vec<i32>> = vec![];
        let v1 = product(n-1, A);
        for &e1 in A {
            for e2 in v1.iter() {
                let e = [&[e1], e2.as_slice()].concat();
                v.push(e)
            }
        }
        v
    }
    else {
        let mut v: Vec<Vec<i32>> = vec![];
        let v1 = product(n/2, A);
        for e1 in v1.iter() {
            for e2 in v1.iter() {
                let e = [e1.as_slice(), e2.as_slice()].concat();
                v.push(e)
            }
        }
        v
    }
}

fn compute_probabilities(n: usize, A: &Vec<i32>) -> Probs {
    let mut probs: Probs = Probs::new();
    if n == 0 {
        probs.insert((0, 0, 0, 0, 0), 1.0);
        return probs
    }
    
    let vs = product(n, A);
    for v in &vs {
        let mut w: [i32; 5] = [0; 5];
        for i in 0..n {
            w[i+5-n] = v[i]
        }
        w.sort();
        let s = array_to_state(&w);
        let e = probs.entry(s).or_insert(0.0);
        (*e) += 1.0 / (vs.len() as f64)
    }
    probs
}

fn score(s: &State) -> i32 {
    let mut counter: HashMap<i32, i32> = HashMap::new();
    for X in [s.0, s.1, s.2, s.3, s.4] {
        let e = counter.entry(X.abs()).or_insert(0);
        (*e) += 1
    }
    counter.into_iter().map(|(X, n)| X * n).max().unwrap()
}

// (-3, -2, 0, 0, 3) + (0, 0, 0, 1, 5) -> (-3, -3, -2, 1, 5)
fn add(s1: &State, s2: &State) -> State {
    let mut v: [i32; 5] = [0; 5];
    let a1 = state_to_array(s1);
    let a2 = state_to_array(s2);
    let mut j: usize = 0;
    for i in 0..5 {
        if a1[i] != 0 {
            v[j] = -(a1[i].abs());
            j += 1
        }
        if a2[i] != 0 {
            v[j] = a2[i];
            j += 1
        }
    }
    v.sort();
    array_to_state(&v)
}

fn mask(s: &State, flags: usize) -> State {
    let mut a = state_to_array(s);
    for i in 0..5 {
        if ((flags >> i) & 1) == 1 {
            a[i] = 0
        }
    }
    a.sort();
    array_to_state(&a)
}

fn is_valid_mask(flags: usize, a: &[i32; 5]) -> bool {
    (0..5).all(|i| ((flags >> i) & 1) == 0 || a[i] > 0)
}

// 最も適切なマスクのときの期待値
fn max_exp_value(s: &State, pss: &Vec<Probs>,
                        E: &HashMap<State, f64>) -> f64 {
    let a = state_to_array(s);
    let mut e_max: f64 = 0.0;
    for flags in 0..32 {
        if !is_valid_mask(flags, &a) {
            continue
        }
        
        let n: usize = (0..5).map(|i| (flags >> i) & 1).sum::<usize>();
        let s1 = mask(s, flags);
        let mut e_each: f64 = 0.0;
        for (&s2, &prob) in pss[n].iter() {
            let s3 = add(&s1, &s2);
            if let Some(&e) = E.get(&s3) {
                e_each += prob * e
            }
        }
        if e_each > e_max {
            e_max = e_each
        }
    }
    e_max
}

fn collect_states(A: &Vec<i32>) -> Vec<State> {
    let mut B: Vec<i32> = A.clone();
    for &a in A {
        B.push(-a)
    }
    
    let vs = product(5, &B);
    let mut set: HashSet<State> = HashSet::new();
    for v in vs {
        let mut a: [i32; 5] = [v[0], v[1], v[2], v[3], v[4]];
        a.sort();
        set.insert(array_to_state(&a));
    }
    set.into_iter().collect::<Vec<_>>()
}

fn F(A: Vec<i32>) -> f64 {
    let mut E: [HashMap<State, f64>; 3] = [HashMap::new(),
                                            HashMap::new(),
                                            HashMap::new()];
    let pss: Vec<Probs> = (0..6).map(|n| compute_probabilities(n, &A)).
                                                                collect();
    let all_states = collect_states(&A);
    for state in all_states.iter() {
        E[0].insert(*state, score(state) as f64);
    }
    
    for state in all_states.iter() {
        E[1].insert(*state, max_exp_value(state, &pss, &E[0]));
    }
    
    for (state, _) in pss[5].iter() {
        E[2].insert(*state, max_exp_value(state, &pss, &E[1]));
    }
    
    pss[5].iter().map(|(&s, &p)| p * E[2].get(&s).unwrap()).sum::<f64>()
}

fn main() {
    let A: Vec<i32> = read_vec();
    println!("{}", F(A))
}