AtCoder Beginner Contest 383 F

https://atcoder.jp/contests/abc383/tasks/abc383_f

単純なDPです。商品をCで並び替えて、合計金額と今の色をすでに使ったかを状態にします。値は満足度の最大値です。今の色の商品が終わったら、全てまだ色を使っていないことにします。

// Diversity
#![allow(non_snake_case)]

use std::cmp::max;


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


//////////////////// DP ////////////////////

const INF: i64 = -1e18 as i64;

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

fn update_dp(dp: &mut Vec<i64>, X: usize, K: i64, prod: Product) {
    let (P, U, _C) = prod;
    for i in (0..dp.len()).rev() {
        let j = i >> 1;
        let k = i & 1;
        if j + P <= X {
            let new_value = dp[i] + U + (if k == 0 { K } else { 0 });
            dp[(j+P)*2+1] = max(dp[(j+P)*2+1], new_value)
        }
    }
}

fn change_color(dp: &mut Vec<i64>) {
    let L = dp.len();
    for i in (0..L).step_by(2) {
        dp[i] = max(dp[i], dp[i+1]);
        dp[i+1] = INF
    }
}


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

type Product = (usize, i64, i32);   // (P, U, C)

fn read_product() -> Product {
    let v: Vec<usize> = read_vec();
    let (P, U, C) = (v[0], v[1] as i64, v[2] as i32);
    (P, U, C)
}

fn read_input() -> (usize, i64, Vec<Product>) {
    let v: Vec<usize> = read_vec();
    let (N, X, K) = (v[0], v[1], v[2] as i64);
    let products: Vec<Product> = (0..N).map(|_| read_product()).collect();
    (X, K, products)
}

fn F(X: usize, K: i64, mut products: Vec<Product>) -> i64 {
    products.sort_by_key(|p| p.2);
    let mut dp = initialize_dp(X);
    let mut prev_C: i32 = 0;
    for prod in products.into_iter() {
        let C = prod.2;
        if C != prev_C {
            change_color(&mut dp);
            prev_C = C
        }
        update_dp(&mut dp, X, K, prod)
    }
    dp.into_iter().max().unwrap()
}

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