Official

F - Diversity Editorial by MtSaka


以下では商品 \(i\)\(i\) 個目の商品を指すとします。

この問題は動的計画法を用いることで解くことができます。具体的には、

\(dp[c][p]=\) 色が \(c\) 以下の商品から \(p\) 円分購入したときの満足度の最大値

として、最終的な答えは \(dp[N][0],dp[N][1],\ldots,dp[N][X]\) の最大値となるようなDPを考えます。

このDPを色ごとにまとめて更新することを考えます。\(dp[c-1]\) まで計算が終わっていて \(dp[c]\) の計算するとき、色が \(c\) である商品を\(1\) 個ずつ見て以下のように \(dp[c]\) の更新をします。商品が \(j\) であるとき、\(p=X,X-1,\ldots,P_j\) の順番で

\(dp[c][p]=\max( \{dp[c-1][p-P_j]+U_j+K,dp[c][p-P_j]+U_j,dp[c][p]\})\)

と更新することで色 \(c\) の商品全てについて更新した後は

\(dp[c][i]=\) (色が \(c\) 以下の商品から 色 \(c\) の商品を \(1\) 個以上購入して \(i\) 円分購入したときの満足度の最大値)

となっています。最後に、各 \(p\) に対して \(dp[c][p]=\max (dp[c][i],dp[c-1][p])\) と更新することで色が \(c\) 以下の商品から \(p\) 円分購入したときの満足度の最大値が \(dp[c][p]\) に格納されます。

まとめると、以下の要領でDPテーブルが計算できます

  • はじめ、\(dp[0][0]=0\) でそれ以外は \(- \infty\) で初期化されているとします。
  • \(c=1,2,\ldots,N\) の順で以下を行います
    • \(c\) の商品 \(i_1,i_2,\ldots,i_k\) であるとき、 \(j=1,2,\ldots,k\) の順番で以下の更新をする
      • \(p=X,X-1,\ldots,P_{i_j}\)の順番で \(dp[c][p]=\max(\{dp[c-1][p-P_{i_j}]+U_{i_j}+K,dp[c][p-P_{i_j}]+U_{i_j},dp[c][p]\})\) と更新する
    • \(p=0,1,\ldots,X\) に順で、\(dp[c][p]=\max (dp[c][p],dp[c-1][p])\) と更新する

商品の個数も色の種類数も \(N\) で更新には \(\mathrm{O}(X)\) かかるため、全体で計算量は \(\mathrm{O}(NX)\) となり、十分高速です。

実装例(C++):

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
static constexpr ll inf = 2e18;

int main() {
    int n, x;
    ll k;
    cin >> n >> x >> k;
    vector<int> p(n);
    vector<ll> u(n), c(n);
    for (int i = 0; i < n; ++i) {
        cin >> p[i] >> u[i] >> c[i];
    }
    vector<vector<int>> id(n);
    for (int i = 0; i < n; ++i) {
        id[c[i] - 1].push_back(i);
    }
    vector<vector<ll>> dp(n + 1, vector<ll>(x + 1, -inf));
    dp[0][0] = 0;
    for (int i = 0; i < n; ++i) {
        for (auto& idx : id[i]) {
            for (int j = x; j >= p[idx]; --j) {
                if (dp[i][j - p[idx]] != -inf)
                    dp[i + 1][j] = max(dp[i + 1][j], dp[i][j - p[idx]] + u[idx] + k);
                if (dp[i + 1][j - p[idx]] != -inf)
                    dp[i + 1][j] = max(dp[i + 1][j], dp[i + 1][j - p[idx]] + u[idx]);
            }
        }
        for (int j = 0; j <= x; ++j)
            dp[i + 1][j] = max(dp[i + 1][j], dp[i][j]);
    }
    cout << *max_element(dp[n].begin(), dp[n].end()) << endl;
}

posted:
last update: