butch’s blog

メモ置き場。

【蟻本】 ハードルが上がった「くじびき」をDPで解く

蟻本(プログラミングコンテストチャレンジブック)のp25に記載されているハードルが上がった「くじびき」をDP(動的計画法)で解きました。 (間違っていたら教えてください。)
問題文はp8に記載されています。

次のDPテーブルを定義します。

  • DP[i][j]:i枚の紙の数字の合計がjになる場合True, それ以外はFalse

iは0~4までの長さ5で設定し、jは紙に書かれた数字の最大値*4を設定します。
初期値はDP[0][0]=True以外は全てFalseにします。
紙に書かれた数字jを入力する時にDP[1][j]=Trueを設定します。
i(>=2)枚以上の紙を使用して和がjになるかどうかを調べていきます。
現在の枚数iよりも1枚少ない状態でDP[i-1][j]=Trueの場合、各紙lを調べてDP[i][j+k[l]] = 1と計算します。
計算量は$O(4nk)$になると思います。
コードは以下になります。

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep2(i, s, n) for (int i = s; i < (int)(n); i++)

const int MAX_N = pow(10, 6);

int main(){
    int n, m;
    cin >> n >> m;
    int k[MAX_N];
    vector<vector<int>> DP(5, vector<int>(MAX_N, 0));
    DP[0][0] = 1;
    int k_max = 0;
    rep(i, n) {
        cin >> k[i];
        DP[1][k[i]] = 1;
        k_max = max(k_max, k[i]);
    }
    k_max *= 4;
    rep2(i, 2, 5) {
        rep(j, k_max) {
            rep(l, n) {
                if (DP[i-1][j]) {
                    DP[i][j+k[l]] = 1;
                }
            }
        }
    }
    if (DP[4][m]) {
        cout << "Yes" << endl;
    } else {
        cout << "No" << endl;
    }
}