問題概要
正整数 $N, K$ が与えられる．次の条件を満たす，有理数の多重集合 $S$ の総数 $\pmod{ 998{,}244{,}353 }$ で求めよ．
- $|S| = N$
- $\sum_{ x \in S } x = K$
- $x \in S$ は $\frac 1 { 2^i }$ ($i = 0, 1, 2, \dots$) の形で書ける
制約
- $1\leq N, K \leq 3{,}000$
解法
$1$ を使う個数 $a$ を決め打ちするとします．このとき，残りの問題は $\frac 1 2$ 以下の有理数 $N - a$ 個で $K - a$ を作る問題が残ります．これは両辺 $2$ 倍すると $1$ 以下の有理数 $N - a$ 個で $2 ( K - a )$ を作る問題だと思うことができて，元々の問題と同じ構造が出てきます．なので，再帰的に問題を解けそうな気持ちになってきます．
答えを求める関数を $f( *, * )$ とします．具体的には，
- $f( n, k ) := \text{$n$ 個の要素で $k$ を作る方法の総数}$
として，中身を設計していきます．まず，明らかに値が定まるケースは 3 つあって，
- $n < k$ のとき，選択肢の中で値が最大の $1$ を最大限使っても足りないので $0$ 通り（invalid な状態）
- $0 < n$ かつ $k = 0$ のとき，有理数が余ってしまっているので $0$ 通り
- $n = 0$ かつ $k = 0$ のとき，出来上がっていてこれ以上何もできないので $1$ 通り
となります．これが再帰のベースケースです．次に再帰ケースについて考えますが．冒頭のように $1$ を使う個数 $a$ を全通り試してしまうとこの部分に $\Theta( n )$ 時間かかってうれしくない（後述しますが TLE します）ので，もう少し考える必要があります．しかし，この部分のアイデアは個数制限無しナップサック問題（蟻本 [1] など）のような感じで，$f( n, k )$ からは $f( n, 2 k )$ と $f( n - 1, k - 1 )$ だけを――$1$ を多重集合に 1 つ加えるかどうかだけを――試せばよいです．これで，例えば $f( n - 2, * )$ は $f( n - 1, * )$ から呼ばれるといった具合に，再帰関数自身がよしなにやってくれます．
$f$ の全体をまとめると，
\begin{align*}
f( n, k ) = \begin{cases}
0 & \text{($n < k$)} \\
0 & \text{($0 < n$ かつ $k = 0$)} \\
1 & \text{($n = 0$ かつ $k = 0$)} \\
f( n, 2k ) + f( n - 1, k - 1 ) & \text{(otherwise)}
\end{cases}
\end{align*}
となります．
$f$ の引数の組合せの総数は $\Theta( N^2 )$ 通りなので，後はこれをメモ化すると TLE せずに問題を解くことができます．途中で書いたような遷移に $\Theta( N )$ 時間かける場合は全体が $\Theta( N^3 )$ 時間となって TLE になります．
参考文献
[1] 秋葉拓哉, 岩田陽一, & 北川宜稔. (2010). プログラミングコンテストチャレンジブック.
コード
constexpr int MOD = 998244353; int dp[ 4096 ][ 1 << 13 ]; int solve( const int N, const int K ) { if ( N < K || ( N && K == 0 ) ) { return 0; } if ( N == 0 && K == 0 ) { return 1; } int &res = dp[N][K]; if ( res != -1 ) { return res; } return res = ( solve( N, 2 * K ) + solve( N - 1, K - 1 ) ) % MOD; } int main() { IN( int, N, K ); fill( AALL( dp ), -1 ); cout << solve( N, K ) << endl; return 0; }
