C++で自動メモ化(競プロ用)

DP問題などを解く際、私はとりあえず再帰関数で書いてみて大丈夫そうならメ モ化再帰に書き換える、という手順を踏むことが多いのですが、この書き換え が若干面倒なので、手間を最小限に抑える方法を考えてみました。

再帰関数はlambdaで書くものとします。lambdaで再帰する方法については C++のラムダで再帰する - koturnの日記 を参照。

再帰関数をメモ化する場合、中身はだいたいこんな感じになるはずです:

int rec(int i, int j) {
    static bool done[N][M] {};
    static int  memo[N][M];

    if(!done[i][j]) {
        int res;
        // res を計算
        memo[i][j] = res;
        done[i][j] = true;
    }
    return memo[i][j];
}

このコードを自動生成することを考えます。

まず、done, memo は多次元配列なので、std::array を用いて多次元配 列型を実装します:

template<typename T, size_t N, size_t... NS>
struct ArrayType {
    using type = array<typename ArrayType<T,NS...>::type,N>;
};

template<typename T, size_t N>
struct ArrayType<T,N> {
    using type = array<T,N>;
};

template<typename T, size_t... NS>
using Array = typename ArrayType<T,NS...>::type;

template<typename T, size_t N>
T& array_at(Array<T,N>& ary, int i) {
    return ary[i];
}

template<typename T, size_t N, size_t... NS, typename... Args>
T& array_at(Array<T,N,NS...>& ary, int i, Args... args) {
    return array_at<T,NS...>(ary[i], args...);
}

template<typename T, size_t N>
const T& array_at(const Array<T,N>& ary, int i) {
    return ary[i];
}

template<typename T, size_t N, size_t... NS, typename... Args>
const T& array_at(const Array<T,N,NS...>& ary, int i, Args... args) {
    return array_at<T,NS...>(ary[i], args...);
}

多次元配列型があれば、以下のようなテンプレートを用いて先ほどのコードを 自動生成できるはずです:

template<typename F, size_t... NS>
class FixPointMemo {
public:
    explicit FixPointMemo(F&& f) : f_(forward<F>(f)) {}

    template<typename... Args>
    decltype(auto) operator()(Args... args) const {
        using R = decltype(f_(*this,args...));
        static Array<bool,NS...> done {};
        static Array<R,NS...>    memo;

        if(!array_at<bool,NS...>(done,args...)) {
            array_at<R,NS...>(memo,args...) = f_(*this,args...);
            array_at<bool,NS...>(done,args...) = true;
        }
        return array_at<R,NS...>(memo,args...);
    }

private:
    F f_;
};

template<size_t... NS, typename F>
decltype(auto) FIXMEMO(F&& f) {
    return FixPointMemo<F,NS...>(forward<F>(f));
}

以上のコードを用意しておけば、このような再帰関数があるとき:

auto f = FIX([&](auto&& self, int i, int j) -> int { /* ... */ });

以下のようにテンプレート引数を追加するだけでメモ化が完了します:

// i,j は [0,1000) の値をとるものとする
auto f = FIXMEMO<1000,1000>([&](auto&& self, int i, int j) -> int { /* ... */ });

関数本体は一切変更しなくてよいためかなり楽だと思います。

難点を挙げると:

• メモ化完了フラグを bool 配列で持っているためメモリ効率が若干悪い (手で書く場合、問題によってはメモ化配列自体に特殊な値を入れることで 代用できる)
• 関数の引数は整数型に限定される

といったところですが、高難易度問題でなければあまり気にしなくていいよう に思います(どうせ工夫を要するDP問題はループで書かざるを得なかったりし ますし)。

この方法を用いてナップサック問題を解いたコードを貼っておきます: 提出 #6436875 - Educational DP Contest / DP まとめコンテスト