【AtCoder解説】動的計画法を使わずにABC178のD - Redistributionを解いてみる

　久しぶりにAtCoderに参加しました。D問題の公式解説が動的計画法を使っていて難しかったので、自分なりの解法をメモしておきたいと思います。用いるのは高校で習う数学Aの知識です。まずは問題文のおさらいです。 

問題文：整数 $S$ が与えられます。 すべての項が $3$ 以上の整数で、その総和が $S$ であるような数列がいくつあるか求めてください。ただし、答えは非常に大きくなる可能性があるので、 ${10}^9+7$ で割った余りを出力してください。

制約 ：1 ≤ S ≤ 2000, 入力はすべて整数

　例として、 S = 11 のときを考えてみましょう。まず数列の最大の項数を考えます。すべての項が３以上という条件から11 ÷ 3 = 3 ･･･ 2 で最大の項数は３であることがわかります。それでは項数が３のとき、S = 11となる数列はいくつ存在するでしょうか。ここでは、玉と箱を用いて考えます。今回はS = 11かつ項数３なので玉を11個、箱を３つ用意しました。箱の中に入っている玉の数が数列の各項に対応します。まず箱には少なくとも３つの玉が入っている必要があるので、あらかじめ３つずつ玉を入れておきます。すると、残った２つの玉を３つの箱に振り分ける方法は何通りあるか、という問題に帰着します。たとえば前２つの箱に１個ずつ入れると、これは数列｛4, 4, 3｝に対応しています。数列｛4, 4, 3｝は確かにS = 11かつ項数３の数列ですね。

　さて、このような振り分け方は何通りあるでしょうか。これを求めるのが「重複組み合わせ」です。残り２つの玉の３つの箱への分け方は、２つの玉と２つの仕切りの並べ方と対応します。この場合は６通りあることがわかります。あとはこれを箱が１つのとき、２つのときの場合の箱への入れ方も計算して、すべてを足せばいいわけです。

解法をまとめます。

1. 入力Sに対して最大の箱の数を求める。
2. 最大の箱の数をXとする。箱の数が１からXのときそれぞれに対して以下の計算を行う。
   ・あらかじめ箱に入れておく玉を入力Sからひく。
   ・残りの玉の箱への入れ方が何通りあるか求める。
   
3. 箱の数が１のときからXのときまでの残りの玉の箱への入れ方を合計する。答えを${10}^9+7$ で割る。

　Pythonでの解答例を示します。今回は階乗の計算にライブラリ「math」の「factorial」wを使いました。階乗の計算ではオーバーフローに注意してください。// のところを　/ とするとfloat型の計算となりオーバーフローします。また、(factorial(ball) * factorial(box - 1)) を factorial(ball) // factorial(box - 1) とするのもオーバーフローの可能性があるようです。

from math import factorial
 
S = int(input())
 
ans = 0
for i in range(S // 3):
    box = i + 1  # 箱の数
    ball = S - 3 * box # あらかじめ箱に３こずつ入れたあとの残りの玉の数
   
    # 箱への入れ方が何通りあるか計算。
    # 仕切りの数は box - 1 であることに注意
    ans_ = factorial(ball + box - 1) // (factorial(ball) * factorial(box - 1))
    ans += ans_
print(ans % (10 ** 9 + 7))

　久しぶりのAtCoder、とても楽しかったです。次も頑張りたいなと思います。