【AtCoder】ABC 390 E - Vitamin Balance | 緑コーダーが解くAtCoder

atcoder.jp

配点: 450 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 1227

問題概要
制約
考察
- DP テーブルの構築
- 二分探索
実装例
コメント

問題概要

$N$ 個の食べ物があり、 $i$ 個目の食べ物を食べると、ビタミン $V_i \: (V_i = 1, 2, 3 のいずれか)$ が $A_i$ 、カロリーが $C_i$ だけ摂取できる。

摂取するカロリーが合計で $X$ 以下となるように0個以上の食べ物を選んで食べるとき、「ビタミン1,2,3のうちもっとも摂取量が少ないものの摂取量」としてあり得る最大値を求めよ。

制約

$1\leq N, X \leq 5000$
$1\leq A_i\leq 2\times 10^5$
$1\leq C_i\leq X$
入力はすべて整数。

考察

合計摂取カロリー $X$ 以下でビタミンの摂取量を最大化する
3種類のビタミン摂取量の最小値を最大化する

という目標について、前者からは「ナップザック DP」、後者からは「二分探索」という方針が想起できる。

DP テーブルの構築

各ビタミンに対して、

$\mathrm{dp}_{i, c} :=$ $i$ 個目までの食べ物について、摂取カロリーが $c$ 以下となるように選んで食べたときの、ビタミン摂取量の最大値

という DP テーブルを考える。この「食べ物」は、今考えているビタミンを摂取できる食べ物のみを指す。これらの食べ物は、「ビタミン摂取量」と「カロリー摂取量」をペアにして配列に格納しておくことにする。

初期値は全て $0$ としておく。

遷移方法は、 $\mathrm{dp}_i$ が求まっていると仮定したとき、 $i$ 番目の食べ物のビタミン摂取量を $\mathrm{vitamin}$ 、カロリー摂取量を $\mathrm{calorie}$ とすると、 $c = 0, 1, \cdots, X$ に対して、

$\mathrm{dp}_{i+1, c} = \left\{ \begin{array}{ll} \max (\mathrm{dp}_{i, c}, \mathrm{dp}_{i, c - \mathrm{calorie}} + \mathrm{vitamin}) & (c \geq \mathrm{calorie}) \quad (\text{食べない or 食べる})\\ \mathrm{dp}_{i, c} & (otherwise) \quad (\text{これ以上は食べられない}) \end{array} \right.$

となる。

欲しいのは、 $\mathrm{dp}_{N}$ という一次元配列である。

なお、公式解説のように、3種類のビタミンをまとめて3次元のテーブルを考えてもよい。

二分探索

前項で得られた一次元配列は、

$\mathrm{dp}_c :=$ カロリーが $c$ 以下で摂取できるビタミンの最大量

であった。

ここで実現したいのは、

あるビタミン摂取量 $m$ について、その量だけビタミンを摂取するために必要な最小のカロリーを各ビタミンごとに計算し、それらの和が $X$ を超えない範囲で $m$ を最大化する

ことである。

少々複雑だが、要は様々な $m$ に対して、条件を満たす範囲と満たさない範囲の境界を探したいということだ。

これは、いわゆる「めぐる式二分探索」によって実現することができる。

$m$ の範囲は $0$ から $\text{(各ビタミンの合計摂取量の最小値)} + 1$ までとなる。

ビタミンを $m$ 摂取するのに必要な最小カロリーを求める部分は、愚直にforループを回せばよい。

計算量は、 DP テーブルの構築に $O(NX)$ 、二分探索部分に $O(X \log M)$ (ただし、 $M$ は各ビタミンの摂取量の最小値) であり、全体として $O(NX + X \log M)$ となる。

実装例

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

using ll = long long;
using Pair_int = pair<int, int>;

constexpr int INF = 1e+9;

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
#define repe(i, start, end) for (auto i = (start); (i) <= (end); (i)++)

// ======================================== //

vector<int> buildDP(const vector<Pair_int> &v, int X)
{
    int n = v.size();
    vector<vector<int>> dp(n + 1, vector<int>(X + 1, 0));
    rep(i, 0, n)
    {
        int vitamin = v[i].first, calorie = v[i].second;
        repe(c, 0, X)
        {
            if (c >= calorie)
            {
                dp[i + 1][c] = max(dp[i][c], dp[i][c - calorie] + vitamin);
            }
            else
            {
                dp[i + 1][c] = dp[i][c];
            }
        }
    }

    return dp[n];
}

int getMinCalorie(const vector<int> &dp, int vitamin)
{
    rep(c, 0, dp.size())
    {
        if (dp[c] >= vitamin)
        {
            return c;
        }
    }

    return INF;
}

int main()
{
    int N, X;
    cin >> N >> X;
    vector<Pair_int> v1, v2, v3;
    ll v1_sum = 0, v2_sum = 0, v3_sum = 0;
    rep(i, 0, N)
    {
        int V, A, C;
        cin >> V >> A >> C;

        if (V == 1)
        {
            v1.push_back({A, C});
            v1_sum += A;
        }
        else if (V == 2)
        {
            v2.push_back({A, C});
            v2_sum += A;
        }
        else
        {
            v3.push_back({A, C});
            v3_sum += A;
        }
    }

    vector<int> dp1 = buildDP(v1, X);
    vector<int> dp2 = buildDP(v2, X);
    vector<int> dp3 = buildDP(v3, X);

    ll ok = 0, ng = min({v1_sum, v2_sum, v3_sum}) + 1;
    while (ng - ok > 1)
    {
        ll mid = (ok + ng) / 2;
        int c1 = getMinCalorie(dp1, mid);
        int c2 = getMinCalorie(dp2, mid);
        int c3 = getMinCalorie(dp3, mid);

        if ((ll)c1 + c2 + c3 <= X)
        {
            ok = mid;
        }
        else
        {
            ng = mid;
        }
    }

    cout << ok << endl;

    return 0;
}