ABC215-G 解いてみた

　組み合わせ問題を多項式の問題に解釈して解くのが好きなので、今回AtCoderBeginnerContest215のG問題を多項式の問題に解釈してみようと思います。(数学の記事のつもりなのでプログラムの実装に関する部分は省略して考察だけしていこうと思います。)

　種類数の期待値は種類数の総和を選び方の総数$\binom{N}{K}$で割ると求まります。種類数の総和を求めていこうと思います。
　$[y^mx^n]f(x,y)$=「$N$個の中から$n$個選んだ時色の種類数が$m$種類になる選び方の総数」となるような多項式を考えていきます。$N$個の中に色$i$のキャンディーが$a_i$個あるとします。結論だけ先に言うと、$f(x,y)$は以下のようになります。
　$f(x,y)=\displaystyle\prod_{i}(1+y((1+x)^{a_i}-1))$
　例えば,$N=6$個の中に赤、青、黄がそれぞれ$1,2,3$個あるとすると,$f(x,y)=(1+yx)(1+2yx+yx^2)(1+3yx+3yx^2+yx^3)$です。
　色$i$ごとに$a_i$個あるうち何個選ぶかを独立に決めてきます。上の例で,黄色キャンディ$3$個のうち$0,1,2,3$個選ぶとするとそれぞれ選び方は$1,3,3,1$通りであり、種類数($y$の指数)への寄与は$0,1,1,1$であり,個数($x$の指数)の寄与はそれぞれ$0,1,2,3$であるので,$(1+3yx+3yx^2+yx^3)$という多項式が思いつきます。すべての色について同様に多項式を考え、積をとると$f(x,y)$となります。
　あとはここから選ぶ個数($x$の指数)を固定したときの種類数($y$の指数)の総和を求めていきます。$y$について微分すると、$y$の指数が係数としておりてくるので、求めたいのは$[x^K]f_y(x,1)$です。
$$\begin{aligned} \quad f_y(x,y)&=\partial_y\prod_{i}(1+y((1+x)^{a_i}-1))\\ &=\sum_i\Big(\partial_y(1+y((1+x)^{a_i}-1))\times\prod_{j\neq i}(1+y((1+x)^{a_j}-1))\Big)\\ &=\sum_i\Big(((1+x)^{a_i}-1)\times\prod_{j\neq i}(1+y((1+x)^{a_j}-1))\Big)\\ \end{aligned}$$
よって、種類数の総和は
$$\begin{aligned} \quad[x^K]f_y(x,1)&=\sum_i[x^K]\Big(((1+x)^{a_i}-1)\times\prod_{j\neq i}(1+x)^{a_j}\Big)\\ &=\sum_i\Big([x^K](1+x)^N-[x^K](1+x)^{N-a_i}\Big)\\ &=\sum_i\Bigg\{\binom{N}{K}-\binom{N-a_i}{K}\Bigg\} \end{aligned}$$
となります。