ITエンジニアのブログ

IT企業でエンジニアやってる人間の日常について

ARC017 [C]無駄なものが嫌いな人

N 個の数値のうち、好きなだけ選んで合計を X にできるパターンの個数を求めるという問題で、 N が小さい値ならばそのまま全探索(O(2^N))すれば良いんですが、 N が 32 だと時間切れになります。

解説を読むと、N 個の数値を半分に分割し、片方の合計を p, もう片方の合計を X-p にできるパターンの個数同士を掛け合わせれば良いようです。考えてみれば確かに、どれを足しあわせるのかを考えなければこの方法が適用できます。

N 個の数値を好きなだけ選んで足し合わせるのは、それぞれ配列の中身を順番に見ていって、それを足した場合と足していない場合について探索を繰り返す事でできるので、思いの外記述量は少なくて済みます。

void search(int a[], int index, int current){
  if(i == n){
    current が合計値
    return;
  }
  search(a, index + 1, current + a[index]); // a[index] を足した場合
  search(a, index + 1, current); // a[index] を足さない場合
}

インデックスを逆にするとグローバル変数や関数の引数を減らせます。

#include <iostream>
#include <map>
#include <iterator>
using namespace std;

void dfs(int index, int current, int *w, map<int, int> &m){
    if(index < 0){
        if(m.count(current)){
            m[current] += 1;
        }else{
            m[current] = 1;
        }
        return;
    }

    dfs(index - 1, current + w[index], w, m);
    dfs(index - 1, current, w, m);
}

int main(){
    int n, x;
    int *wf, *ws;
    map<int, int> mf;
    map<int, int> ms;
    int result = 0;

    cin >> n >> x;

    if(n == 1){
        int w;
        cin >> w;
        cout << (w == x ? 1 : 0) << endl;
        return 0;
    }

    int nf = n / 2;
    int ns = n - n / 2;

    wf = new int[nf];
    ws = new int[ns];

    for(int i = 0; i < nf; i++){
        cin >> wf[i];
    }

    for(int i = 0; i < ns; i++){
        cin >> ws[i];
    }

    dfs(nf - 1, 0, wf, mf);
    dfs(ns - 1, 0, ws, ms);

    for(map<int, int>::iterator it = mf.begin(); it != mf.end(); it++){
        if(ms.count(x - it->first)){
            result += it->second * ms[x - it->first];
        }
    }

    cout << result << endl;

    delete [] wf;
    delete [] ws;

    return 0;
}

マップには値 p となる組み合わせが patterns 通りとして map[p] = patterns となっているので、それぞれに対してもう片方の map[X - p] と掛けあわせます。その合計が全てのパターンになります。

map の iterator とか初めて使った。 Python だと for key in dict で全てのキーを拾えるので楽です。しかし iterator も結構簡潔に書けますね。