CodeChef May Challenge 2020 - Chef and Rainbow Road

• 問題
• AC

要約: $a_1, ..., a_N$から (1) 重複を許して選んで (2) どれも少なくとも1個以上選んで (3) 合わせて$N+K-1$個になるように選んで、積を取るとする。それらの総和を求めよ。

最低$1$個以上という条件を外して$K-1$個選ぶことにして、求まった値に$a_1...a_N$を掛ければよい。

多項式で考えるんだろうなという検討はつくが、ほとんど何も知らないのでmaspyさんの記事などを見て式を考える。とりあえず

$(1+a_1x+a_1^2x^2 + ...)(1+a_2x+a_2^2x^2 + ...) ... (1+a_Nx+a_N^2x^2 + ...)$

の第$K-1$項(0-based)を見ればよいことがわかる。この式は$((1-a_1x)(1-a_2x)...(1-a_Nx))^{-1}$と表せて、$(1-a_1x)(1-a_2x)...(1-a_Nx)$は分割統治してFFTすればよく、逆元もFFTで求まるらしいので、$K \le 10^5$の部分点1, 2はそれで取れる。

部分点3からはこの方針では無理そう。ところで、この間のlong challengeのまったく歯が立たなかった問題についてNyaanさんが記事を書いていたのを思い出して見に行ったら、部分分数分解というワードが出てきて使えそうと思う。

$$\frac{1}{(1-a_1x)(1-a_2x)} = \frac{a_1}{(a_1-a_2)(1-a_1x)} + \frac{a_2}{(a_2-a_1)(1-a_2x)}$$

と分解できるので、全部掛けると$1/(1-a_1 x)$の係数は$a_1^{N-1}/\prod_{j \neq 1} (a_1 - a_j)$になる。対称性により$1/(1-a_i x)$の係数は$a_i^{N-1}/\prod_{j \neq i} (a_i - a_j)$であり、これが$1+a_i x + a_i^2 x^2 + ...$に掛かっているので$x^{K-1}$の係数は$a_i^{K+N-2}/\prod_{j \neq i} (a_i - a_j)$とわかる。この値をすべての$i$について加算すれば答えが得られる。部分点3が取れた。

満点について。$a_i^{K+N-2}/\prod_{j \neq i} (a_i - a_j)$の分母の形を見るとヴァンデルモンド行列式やラグランジュ補間などが思い浮かぶので、ラグランジュ補間の解説を読んでいて$(x-a_1)...(x-a_N)$を微分したものに代入すればいいと気づく。多項式の値を$N$個の点について求める必要があるが、ググると${O}(N\log^2 N)$でできるらしい → 間に合った。

---

最終日に満点に気づいた。時間がなさすぎたのでbeetさんのライブラリを大急ぎで写経…… 長期コンテストでも終了直前に方針が立つと結局時間に追われることになるな。