Official

C - Dice Sum Editorial by nok0


数列を全通り試そうとすると、 \(M^N\) 種類の数列が考えられるため、実行時間制限には到底間に合いません。

数列を先頭から決めていくこととします。 ここで、以下の事実に着目します。

  • 数列を先頭から決めていく際、覚えておくべき必要があるものはその時点での数列の総和のみであり、具体的に各要素の値が何であったかは捨象してよい。

このことを用いると、動的計画法で解くことができます。

DP の定義

\(\mathrm{dp}[i][j]:=\)数列の先頭から \(i\) 項まで決めた際に、総和が \(j\) であるような数列の決め方の総数

\(i\) としては \(0\) 以上 \(N\) 以下、\(j\) としては \(1\) 以上 \(K\) 以下を考えれば良いので、状態数は \(\mathrm{Ο}(NK)\) です。

DP の初期値

はじめ、 \(\mathrm{dp}[0][0] = 1\) で、残りは \(0\) で埋めます。

DP の遷移

\(i\) の昇順に計算していきます。 \(j = 0, \ldots, K-1\)\(k = 1,\ldots,M\) について

  • もし \(j + k \leq K\) ならば、 \(\mathrm{dp}[i+1][j + k]\)\(\mathrm{dp}[i][j]\) を加算する。

と遷移すればよいです。

答え

最終的に、 \(\mathrm{dp}[N][1] + \mathrm{dp}[N][2] + \ldots + \mathrm{dp}[N][K]\) が答えとなります。

計算量は \(\mathrm{Ο}(NMK)\) です。

実装例 (C++):

#include <bits/stdc++.h>
#include <atcoder/modint>

using namespace std;
using mint = atcoder::modint998244353;

int main() {
	int n, m, K;
	cin >> n >> m >> K;

	vector dp(n + 1, vector(K + 1, mint(0)));
	dp[0][0] = 1;

	for(int i = 0; i < n; i++) {
		for(int j = 0; j < K; j++) {
			for(int k = 1; k <= m; k++) {
				if(j + k <= K) dp[i + 1][j + k] += dp[i][j];
			}
		}
	}

	mint res = 0;
	for(int i = 1; i <= K; i++) {
		res += dp.back()[i];
	}

	cout << res.val() << endl;
}

Bonus 1

\(\mathrm{Ο}(NK)\) で解いてみましょう。(中級者向け)

方針

解説では配る DP でしたが、貰う DP で考えます。

このとき、 $\mathrm{dp}[i + 1][j]$ はある整数 $L,R$ $(L\leq R)$ を用いて

$\mathrm{dp}[i + 1][j] = \mathrm{dp}[i][L] + \ldots + \mathrm{dp}[i][R]$

と書けます。これは区間和の形をしているので、 $\mathrm{dp}[i+1]$ を更新する際に $\mathrm{dp}[i]$ の累積和を前計算しておくことで、遷移が $\mathrm{Ο}(1)$ になります。

以上の工夫により $\mathrm{Ο}(NK)$ でこの問題に答えることができました。累積和を用いた DP の高速化は頻出なので、覚えておくと良いでしょう。

Bonus 2

\(\mathrm{Ο}(K)\) で解いてみましょう。(上級者向け)

方針

形式的冪級数を用いて解きます。

$f(x)=x + x^2 + \ldots + x^ M = x \frac{1-x^M}{1-x}$ とすると、答えは $f(x)^N$ の $1$ 次 から $K$ 次の係数の和です。

ここで、 $f(x)$ の $0$ 次 から $K$ 次の係数の和が $f(x)/(1-x)$ の $K$ 次の係数と一致することを用いて言い換えると、 $(1-x^M)^{N}(1-x)^{-(N+1)}$ の $K-N$ 次の係数が求められれば良いです。

$(1-x^M)^{N}$ の各項の係数は二項定理により $\mathrm{Ο}(K)$ で計算できます。また、$(1-x)^{-(N+1)}$ の各項の係数も負の二項定理(ここでは詳しく解説しません、興味のある方は調べてください)により $\mathrm{Ο}(K)$ で求めることができます。

以上で、この問題を $\mathrm{Ο}(K)$ で解くことができました。形式的冪級数を用いた数え上げの高速化は高難易度の問題で時々見かけるので、興味のある方は是非勉強してみてください。

posted:
last update: