Typical DP Contest H - ナップザック

解法

まず,色でソートしておく.
すると以下の dp が書ける.
dp[i][w][j][k] := i 番目まで見たとき,j 種類の色を使っていて,最後に使った色が k であって,重さの総和が w となるときの最大値
遷移は素直に書くだけなので省略します.
注意することとしてはループの回す方向を考えないといけないところ.

計算量は O(NC^2W).

ソースコード

#include <bits/stdc++.h>
using namespace std;

constexpr int INF = 1e9;

struct baggage {
    int w, v, c;
    bool operator<(baggage const& b) const {
        return c < b.c;
    }
};

int main() {
    int N, W, C;
    cin >> N >> W >> C;
    vector<baggage> b(N);
    for(int i = 0; i < N; ++i) {
        cin >> b[i].w >> b[i].v >> b[i].c;
    }
    sort(begin(b), end(b));
    vector<vector<vector<int>>> dp(W + 1, vector<vector<int>>(C + 1, vector<int>(51, -INF)));
    dp[0][0][0] = 0;
    int res = 0;
    for(int i = 0; i < N; ++i) {
        for(int w = W - b[i].w; w >= 0; --w) {
            for(int cn = C; cn >= 0; --cn) {
                for(int last_c = 0; last_c <= b[i].c; ++last_c) {
                    int next_cn = cn + (last_c != b[i].c);
                    if(dp[w][cn][last_c] == -INF || next_cn > C) {
                        continue;
                    }
                    dp[w + b[i].w][next_cn][b[i].c] = max(dp[w + b[i].w][next_cn][b[i].c], dp[w][cn][last_c] + b[i].v);
                    res = max(res, dp[w + b[i].w][next_cn][b[i].c]);
                }
            }
        }
    }
    cout << res << endl;
}

感想

すんなり解けたけど,多分高速化の余地がある.
定数倍入れても大体 C^2NW = 10^8 なのでちょっと危ない(実際 1sec かかる).
わかったら追記します.