【C++】重みの合計がちょうどXで、重みMの辺を少なくとも1つ含むパスの個数を求める方法

問題概要

無限の階層を持ちうる木構造が与えられます。変数 child は1つのノードが持てる子の数(= 各辺の重みとして選べる値の上限)、変数 weight は必ず1回以上含めたい辺の重み M、変数 path はパス全体の重みの合計 X を表します。この記事では、重みの合計がちょうど X に等しく、かつ重み M を持つ辺を少なくとも1つ含むパスの個数を求める方法を解説します。

入力例1

入力 − int child = 4, weight = 4, path = 4;

出力 − 重みがちょうどXで、重みMの辺を少なくとも1つ含むパスの個数:1

説明 − ノードには4つの子があり、重み4の辺で構成されるパスが考えられます。合計が4になるのは「重み4の辺を1本だけ使う」ケースのみのため、カウントは1となります。

入力例2

入力 − int child = 3, weight = 2, path = 4;

出力 − 重みがちょうどXで、重みMの辺を少なくとも1つ含むパスの個数:4

説明 − ノードには3つの子があり、重み2の辺を含む必要があります。合計が4になる組み合わせは「2+2」「2+1+1」「1+2+1」「1+1+2」の4通りで、いずれも重み2の辺を含むため、カウントは4となります。

アルゴリズムのアプローチ

• 子の総数、パスの重みの合計、対象となる辺の重みを、それぞれ変数 child、path、weight に入力します。
• 必要なサイズの二次元配列を宣言します。
• i を 0 から配列サイズまでループし、内側のループで j を 0 以上 2 未満の範囲で回しながら arr[i][j] を -1 に初期化します。-1 は「まだ計算されていない」ことを意味します。
• 関数 total_weight() を path、0、weight、child、arr を引数として呼び出します。第2引数の 0 は「重み M の辺をまだ使用していない」ことを示すフラグです。
• 関数 total_weight() 内部では、次の手順を実行します。
  • 結果を格納する一時変数 count を宣言します。
  • path が 0 未満の場合は 0 を返します。
  • path が 0 の場合はフラグ i をそのまま返します(重み M の辺を既に使っていれば 1、未使用なら 0)。
  • arr[path][i] が -1 以外の場合は、計算済みの値を返します(メモ化による高速化)。
  • j を 1 から child までループします。j が weight と等しい場合は、total_weight(path - j, 1, weight, child, arr) の再帰呼び出し結果を count に加算します(重み M の辺を使ったためフラグを 1 に更新)。
  • それ以外の場合は、total_weight(path - j, i, weight, child, arr) の結果を count に加算します。
  • count の値を arr[path][i] に保存します。
• arr[path][i] を返します。
• 最終的な結果を出力します。

C++ 実装例

#include <bits/stdc++.h>
using namespace std;
#define size 4
#define col 4
int total_weight(int path, int i, int weight, int child, int arr[size + 1][col]) {
    int count = 0;
    if (path < 0) {
        return 0;
    }
    if (path == 0) {
        return i;
    }
    if (arr[path][i] != -1) {
        return arr[path][i];
    }
    for (int j = 1; j <= child; j++) {
        if (j == weight) {
            count += total_weight(path - j, 1, weight, child, arr);
        } else {
            count += total_weight(path - j, i, weight, child, arr);
        }
    }
    arr[path][i] = count;
    return arr[path][i];
}
int main() {
    int child = 4, weight = 4, path = 4;
    int arr[size + 1][col];
    for (int i = 0; i <= size; i++) {
        for (int j = 0; j < 2; j++) {
            arr[i][j] = -1;
        }
    }
    cout << "Count of number of paths whose weight is exactly X and has at-least one edge of weight M are: " << total_weight(path, 0, weight, child, arr);
}

上記のコードを実行すると、次の出力が得られます。

出力

Count of number of paths whose weight is exactly X and has at-least one edge of weight M are: 1

計算量について

メモ化により、同じ状態 (path, i) の計算は一度しか行われません。状態数は path × 2、各状態で child 通りの遷移を試すため、時間計算量は O(path × child)、空間計算量は O(path) となります。素朴な再帰では指数時間かかるケースでも、この手法なら効率的に正解を求められます。