汎用性の高そうな数え上げ問題

某サーバーで出た問題の言い換えが何かに使えそうだったのでメモ

問題

$ N 個の区別できる箱それぞれに、$ 0 ～$ \sigma - 1 個区別できないボールを入れて行って、

ボールの総数が$ K 以下かつ$ =K \pmod \sigma になるパターンは何通りか

解答

ボールの総数は多くても$ (\sigma -1) N なので、$ K は「$ (\sigma -1)N よりちょっと大きくて$ \bmod \sigma で$ K と等しい数」で考えていい。

指数が「使ったボールの総数」となるFPSを考えると、

1つの箱に入れるボールのパターン：$ 1 + x + \dots + x^{\sigma-1} = (1-x^\sigma)/(1-x)

その箱が$ N 個：$ (1-x^\sigma)^N/(1-x)^N

$ \bmod \sigma で考える：↑に$ 1 + x^\sigma + x^{2\sigma}+\dots = 1/(1-x^\sigma) を掛けて、$ (1-x^\sigma)^{N-1}/(1-x)^N

最後の式の$ x^K の係数が答えになる。

$ (1-x^\sigma)^{N-1} は展開したとき、項が$ x^0, x^\sigma, x^{2\sigma}, \dots, x^{(N-1)\sigma} の$ N 個しかないので、

$ (1-x)^{-N} は$ x^K, x^{K-\sigma}, x^{K-2\sigma},\dots, x^{K-(N-1)\sigma} の$ N 項だけ計算すればいい。

前者の$ x^{i\sigma} の係数は二項定理から$ (-1)^i\binom{N-1}{i} 、

後者の$ x^{K-i\sigma} の係数は負の二項定理から$ \binom{K-i\sigma+N-1}{N-1} となるので、

結局、答えは$ \sum_{i=0}^{N-1} (-1)^i\binom{N-1}{i}\binom{K-i\sigma+N-1}{N-1}となる。

二項係数の前計算に$ O(K+N) = O(\sigma N) 、総和の計算に$ O(N) かかるので、全体の計算量は$ O(\sigma N) となる。