Official

F - Diversity Editorial by en_translator


Let us call the \(i\)-th product “product \(i\).”

This problem can be solved with Dynamic Programming (DP). Specifically, consider a DP where

\(dp[c][p]=\) the maximum satisfaction when buying products with colors \(c\) or less for \(p\) yen or less.

The final answer is the maximum value among \(dp[N][0],dp[N][1],\ldots,dp[N][X]\).

We will try to update this DP at once for each color. When we already know \(dp[c-1]\) and want to find \(dp[c]\), we inspect the products of color \(c\) one by one to update \(dp[c]\) as follows. When inspecting product \(j\), we set

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

for each \(p=X,X-1,\ldots,P_j\) in this order. When these updates are done for all products of color \(c\), we have

\(dp[c][i]=\) the maximum satisfaction when buying products with colors \(c\) or less for \(p\) yen or less, where at least one product of color \(c\) is chosen.

Finally, we set \(dp[c][p]=\max (dp[c][i],dp[c-1][p])\), so that \(dp[c][p]\) stores the maximum satisfaction when buying products with colors \(c\) or less for \(p\) yen or less.

In summary, the DP table can be computed as follows.

  • Initialize \(dp[0][0]\) with \(0\), and others with \(- \infty\).
  • For \(c=1,2,\ldots,N\) in order, do the following.
    • When products \(i_1,i_2,\ldots,i_k\) have color \(c\), apply the following update for each \(j=1,2,\ldots,k\) in order.
      • For each \(p=X,X-1,\ldots,P_{i_j}\) in order, set \(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]\})\).
    • For \(p=0,1,\ldots,X\) in order, set \(dp[c][p]=\max (dp[c][p],dp[c-1][p])\).

There are \(N\) products and \(N\) kinds of colors, and each update costs \(\mathrm{O}(X)\), so the overall complexity is \(\mathrm{O}(NX)\), which is fast enough.

Sample code (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: