D. Make Them Equal (Educational Codeforces Round 122 Div. 2)

問題

https://codeforces.com/contest/1633/problem/D

問題概要

長さnですべての要素が1の配列 a が与えられる。配列aの要素に対して、a_i = a_i + a_i / x (x > 0) という操作が最大 k 回実行可能である。

長さnの配列bとcがあり、上記の操作で a_i == b_i とできた場合、c_i枚のコインを獲得することができる。

最大何枚のコインを獲得できるか答えよ。

考察

1から各bまでの距離がわかれば、その距離を重さ、コインを価値としたナップサック問題とみなすことができます。

ここで

• bまでの距離をどうやって効率良く求めるか？
• kが10⁶なので間に合わない？

の2点が問題となります。

bまでの距離をどうやって効率良く求めるか？

bはたかだか1000なので、BFSなりメモ化再帰なりで求めることが可能です。これを事前に計算しておけばよいです。

このパートだけ切り取れば、AtCoder Beginner ContestのCかD問題で出てきそうなレベルの問題かと思います。

kが10⁶なので間に合わない？

前計算の結果を確認すると、1000までのいずれかに対する距離はたかだか12であることがわかります。nは10³なので、距離の合計の最大値は最悪でも12000です。よって、ナップサックDPで解くことが可能です。

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

const int INF = 1001001001;

int main() {

    // 先に1000までの距離を求める
    vector<int> a(1001, INF);

    function<void(int, int)> rec = [&](int m, int u) {
        if (m > 1000) return;
        if (a[m] <= u) return;
        a[m] = u;

        for (int x = 1; x <= m; x++) {
            rec(m + m / x, u + 1);
        }
    };

    rec(1, 0);

    int t;
    cin >> t;
    while (t--) {
        int n, k;
        cin >> n >> k;
        vector<int> b(n);
        vector<int> c(n);
        int mx = 0;
        for (int i = 0; i < n; i++) {
            cin >> b[i];
            mx += a[b[i]];
        }
        int sum = 0;
        for (int i = 0; i < n; i++) {
            cin >> c[i];
            sum += c[i];
        }

        // 距離の合計が k 以下なら全てのコインが取得可能
        if (mx <= k) {
            cout << sum << endl;
            continue;
        }

        vector<vector<int>> dp(n + 1, vector<int>(k + 1, -INF));
        dp[0][0] = 0;
        for (int i = 0; i < n; i++) {
            for (int j = 0; j <= k; j++) {
                if (dp[i][j] == -INF) continue;
                if (dp[i + 1][j] < dp[i][j]) dp[i + 1][j] = dp[i][j];
                if (j + a[b[i]] <= k) {
                    if (dp[i + 1][j + a[b[i]]] < dp[i][j] + c[i]) dp[i + 1][j + a[b[i]]] = dp[i][j] + c[i];
                }
            }
        }
        int ans = 0;
        for (int j = 0; j <= k; j++) {
            if (ans < dp[n][j]) ans = dp[n][j];
        }
        cout << ans << endl;
    }
    return 0;
}