AtCoder Beginner Contest 390 E

https://atcoder.jp/contests/abc390/tasks/abc390_e

各々のビタミンの種類がA以上になるときにカロリーの合計がX以下になるか調べるとよいです。そうすると、ビタミンの種類ごとに考えることができてDPで簡単にA以上になるときのカロリーの最小値が求められます。

// Vitamin Balance
#![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()
}


//////////////////// Food ////////////////////

type Food = (usize, i64, i32);

fn read_food() -> Food {
    let v: Vec<i32> = read_vec();
    let (V, A, C) = (v[0] as usize, v[1] as i64, v[2]);
    (V, A, C)
}


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

fn read_input() -> (i32, Vec<Food>) {
    let v: Vec<i32> = read_vec();
    let (N, X) = (v[0], v[1]);
    let foods: Vec<Food> = (0..N).map(|_| read_food()).collect();
    (X, foods)
}

fn collect_foods_by_vitamin(foods: Vec<Food>) -> [Vec<Food>; 3] {
    let mut foods_table: [Vec<Food>; 3] = Default::default();
    for food in foods.into_iter() {
        foods_table[food.0-1].push(food)
    }
    foods_table
}

use std::cmp::max;

fn initialize_dp(X: i32) -> Vec<i64> {
    let mut dp: Vec<i64> = vec![-1; (X as usize)+1];
    dp[0] = 0;
    dp
}

// 同じカロリーで最大のビタミン
fn update_dp(food: &Food, dp: Vec<i64>) -> Vec<i64> {
    let mut new_dp = dp.clone();
    let C = food.2 as usize;
    for c in C..dp.len() {
        if dp[c-C] != -1 {
            new_dp[c] = max(new_dp[c], dp[c-C] + food.1 as i64)
        }
    }
    new_dp
}

// ビタミンがA以上のときの最小のカロリー
fn F_A_each(A: i64, X: i32, foods: &Vec<Food>) -> i32 {
    let mut dp = initialize_dp(X);
    for food in foods.iter() {
        dp = update_dp(food, dp)
    }
    
    for x in 0..dp.len() {
        if dp[x] >= A {
            return x as i32
        }
    }
    X + 1
}

// 3種類のビタミンの最小がA以上になるか
fn F_A(A: i64, X: i32, foods: &[Vec<Food>; 3]) -> bool {
    let x = (0..3).map(|i| F_A_each(A, X, &foods[i])).sum::<i32>();
    x <= X
}

fn F(X: i32, foods: Vec<Food>) -> i64 {
    let foods_table = collect_foods_by_vitamin(foods);
    
    let mut A = 1;
    while F_A(A, X, &foods_table) {
        A *= 2
    }
    let mut first = A / 2;
    let mut last = A;
    while last - first > 1 {
        let mid = (first + last) / 2;
        if F_A(mid, X, &foods_table) {
            first = mid
        }
        else {
            last = mid
        }
    }
    first
}

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