C++で探索コストがちょうどKとなる配列を構築する方法

問題概要

3つの整数 n、m、k が与えられます。正の整数からなる配列の最大要素を求めるために、次のようなアルゴリズムが用意されています。

max_val := -1
max_ind := -1
search_cost := 0
n := size of arr
for initialize i := 0, when i < n, update (increase i by 1), do:
    if max_val < arr[i], then:
        max_val := arr[i]
        max_ind := i
        (increase search_cost by 1)
return max_ind

このアルゴリズムは配列を先頭から順に走査し、それまでの最大値 max_val よりも大きい要素に出会うたびに最大値を更新し、そのたびに search_cost を1増やします。つまり search_cost とは「走査中に最大値が更新された回数」を意味します。

ここでの課題は、次の条件をすべて満たす配列 arr を構築することです。

• 要素数: arr はちょうど n 個の整数を含む
• 値の範囲: すべての要素 arr[i](0 <= i < n)は 1 以上 m 以下の範囲に収まる
• 探索コスト: 上記アルゴリズムを arr に適用したとき、search_cost の値がちょうど k と等しくなる

この条件を満たす配列の構成方法が何通りあるかを求めてください。答えは非常に大きな値になり得るため、10^9 + 7 で割った余りを出力します。

たとえば入力が n = 2、m = 3、k = 1 のとき、答えは 6 になります。条件を満たす配列は [1, 1]、[2, 1]、[2, 2]、[3, 1]、[3, 2]、[3, 3] の6通り存在するためです。

解き方のアプローチ

この問題は、メモ化再帰(トップダウン方式の動的計画法)を使って効率的に解けます。状態として「現在の位置 idx」「残りの探索コスト k」「それまでに見た最大値 currVal」の3つを持ちます。具体的な手順は以下の通りです。

• MOD := 10^9 + 7 と定義する。
• 2つの値を加算して剰余を取る関数 add(a, b) を用意する。戻り値は ((a mod MOD) + (b mod MOD)) mod MOD。
• サイズ 54 × 54 × 105 の3次元配列 dp を宣言する。
• 再帰関数 help(idx, m, k, currVal, n) を次のように実装する。
  • k < 0 の場合は 0 を返す(コストを使い果たした無効な経路)。
  • idx == n + 1 の場合は、k == 0 なら 1(true)、そうでなければ 0(false)を返す。
  • dp[idx][k][currVal + 1] が -1 以外なら計算済みなので、その値を返す(メモ化)。
  • ret := 0 で初期化し、i を 1 から m までループする。
    • i > currVal の場合(新しい最大値が出現しコストが1増える): ret = add(help(idx + 1, m, k - 1, max(currVal, i), n), ret)
    • それ以外の場合(コストは変化しない): ret = add(help(idx + 1, m, k, max(currVal, i), n), ret)
  • 結果を dp[idx][k][currVal + 1] に保存して返す。
• メイン側では、まず dp 配列の全要素を -1 で初期化し、help(1, m, k, -1, n) の結果を返す。

currVal に 1 を加えてインデックスにしているのは、currVal の初期値が -1 であるため、配列の添字として扱えるようにオフセットしているからです。

実装例(C++)

理解を深めるために、実際のC++による実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
const lli m = 1e9 + 7;
class Solution {
public:
    lli add(lli a, lli b) {
        return ((a % m) + (b % m)) % m;
    }
    int dp[54][54][105];
    int help(int idx, int m, int k, int currVal, int n) {
        if (k < 0)
             return 0;
        if (idx == n + 1) {
             return k == 0;
        }
        if (dp[idx][k][currVal + 1] != -1)
             return dp[idx][k][currVal + 1];
        int ret = 0;
        for (int i = 1; i <= m; i++) {
             if (i > currVal) {
                 ret = add(help(idx + 1, m, k - 1, max(currVal, i), n), ret);
             }
             else {
                 ret = add(help(idx + 1, m, k, max(currVal, i), n), ret);
             }
        }
        return dp[idx][k][currVal + 1] = ret;
    }
    int numOfArrays(int n, int m, int k) {
        for (int i = 0; i < 54; i++)
             for (int j = 0; j < 54; j++)
                 for (int k = 0; k < 105; k++)
                     dp[i][j][k] = -1;
        int ret = help(1, m, k, -1, n);
        return ret;
    }
};
main() {
    Solution ob;
    cout << (ob.numOfArrays(2, 3, 1));
}

入力

2, 3, 1

出力

6

まとめ

本記事では、「最大値の更新回数(search_cost)がちょうど k 回になるような長さ n の配列」の構成方法の総数を、10^9 + 7 で割った余りで求める問題を取り上げました。位置・残りコスト・現時点の最大値を状態としたメモ化再帰により、指数爆発しがちな全列挙を効率的に数え上げることができます。「複数の状態を持ちながら数え上げる」タイプの問題では、今回のような多次元DPの枠組みが広く応用できるので、ぜひマスターしておきましょう。