AtCoder Beginner Contest 404 D

https://atcoder.jp/contests/abc404/tasks/abc404_d

動物園に行く回数は3通りなので合計通りです。計算量は、となり十分に間に合います。状態を動物を何回見たかとしますが、1種類の2ビットで表せるので、u128を2つ使えばよいです。同じ状態で最小の金額だけ保持しておけばよいですが、十分に間に合うのでそうしなくてもよいです。

// Goin' to the Zoo
#![allow(non_snake_case)]

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


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

fn read_input() -> (Vec<i64>, Vec<Vec<usize>>) {
    let v: Vec<usize> = read_vec();
    let M = v[1];
    let C: Vec<i64> = read_vec();
    let mut A: Vec<Vec<usize>> = vec![vec![]; M];
    for i in 0..M {
        let w: Vec<usize> = read_vec();
        A[i] = w.iter().skip(1).map(|&i| i-1).collect::<Vec<_>>()
    }
    (C, A)
}

// 動物園ごとにどの動物がいるかに変換する
fn convert(N: usize, A: Vec<Vec<usize>>) -> Vec<Vec<usize>> {
    let mut B: Vec<Vec<usize>> = vec![vec![]; N];
    for (i, v) in A.into_iter().enumerate() {
        for k in v.into_iter() {
            B[k].push(i)
        }
    }
    B
}

type DP = Vec<(u128, u128, i64)>;

fn update_bit(d: u128, d1: u128) -> u128 {
    let mut d2: u128 = 0;
    for k in 0..64 {
        let q1 = (d >> (k*2)) & 3;
        let q2 = (d1 >> (k*2)) & 3;
        let q = min(2, q1 + q2);
        d2 |= q << (k*2)
    }
    d2
}

fn update_dp(dp: DP, v: &Vec<usize>, C: i64) -> DP {
    let mut new_dp: DP = DP::new();
    for (d, u, c) in dp.into_iter() {
        let mut d1: u128 = 0;
        let mut u1: u128 = 0;
        for &k in v.iter() {
            if k < 64 {
                d1 += 1 << (k * 2)
            }
            else {
                u1 += 1 << ((k - 64) * 2)
            }
        }
        for i in 0u128..3 {
            new_dp.push((update_bit(d, d1*i),
                         update_bit(u, u1*i), c+C*(i as i64)))
        }
    }
    new_dp
}

fn calc_full(M: usize) -> (u128, u128) {
    let mut d: u128 = 0;
    let mut u: u128 = 0;
    for i in 0..M {
        if i < 64 {
            d |= 2 << (i*2)
        }
        else {
            u |= 2 << (i*2)
        }
    }
    (d, u)
}

fn F(C: Vec<i64>, A: Vec<Vec<usize>>) -> i64 {
    let N = C.len();
    let M = A.len();
    let B = convert(N, A);
    let mut dp: Vec<(u128, u128, i64)> = vec![(0, 0, 0)];
    for i in 0..N {
        dp = update_dp(dp, &B[i], C[i])
    }
    let du_full = calc_full(M);
    let mut c: i64 = i64::MAX;
    for (d, u, c1) in dp {
        if (d, u) == du_full {
            c = min(c, c1)
        }
    }
    c
}

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