Typical DP Contest - H ナップザック

H: ナップザック - Typical DP Contest | AtCoder

//
// * 普通に考えると dp[used_color_sum][weight_sum] を更新する。
//  => どの色を使ったのかの判定はどうする？
//  => 色ごとに DP を更新していけばよい。
//
class h_knapsack {
public:
    void solve(void) {
            int N,W,C;
            cin>>N>>W>>C;

            map<int,vector<pair<int,int>>> items;
            REP(i, N)
            {
                int w,v,c;
                cin>>w>>v>>c;
                // 色ごとに振り分ける
                items[c].emplace_back(w,v);
            }

            // dp[今まで使った色の合計][重さの合計]
            vector<vector<int>> dp(C+1, vector<int>(W+1, 0));

            // 色ごとに更新する
            for (const auto &pp : items)
            {
                const auto &citems = pp.second;
                // 現在の色の品物を必ず使う DP
                vector<vector<int>> dp2(dp);
                for (int i = 0; i < citems.size(); ++i)
                {
                    int w, v;
                    tie(w, v) = citems[i];
                    for (int c_sum = C; c_sum >= 0; --c_sum)
                    for (int w_sum = W; w_sum >= w; --w_sum)
                    {
                        // dp2[i+1][c_sum][w_sum+w] = max(. , dp2[i][c_sum][w_sum]+v)
                        // i が 1 刻みなので同じ配列を使いまわす
                        dp2[c_sum][w_sum] = max(dp2[c_sum][w_sum], dp2[c_sum][w_sum-w]+v);
                    }
                }
                // 前回の結果にマージ
                REP(c_sum, C+1)
                REP(w_sum, W+1)
                {
                    if (c_sum+1 <= C) // dp2[c_sum][] は c_sum に現在の色を追加したものなので c_sum+1 と比較
                        dp[c_sum+1][w_sum] = max(dp[c_sum+1][w_sum], dp2[c_sum][w_sum]);
                }
            }
            int ans = 0;
            REP(c_sum, C+1)
            REP(w_sum, W+1)
                ans = max(ans, dp[c_sum][w_sum]);
            cout<<ans<<endl;
    }
};