\(f(x)={}_AC_0+{}_AC_1x+\dots+{}_AC_Nx^N\) とする。\(n \geq 2\) に対し、レベル \(n\) の宝石の種類数を \(a_n\) とおくと、\(a_n=[x^n]f(x)(1+a_2 x) \dots (1+a_{n-1}x)\) となる。

よって、\(i = 2, 3, \dots, N\) の順に以下の操作をすればよい。

• \([x^i]f(x)\) の値を調べる。その答えを \(a_i\) とする。
• \(f(x) \leftarrow f(x) (1+a_i x)\) と更新する。

\(f(x)\) の更新をするとき、\((1+a_i x)\) をそのまま掛けるかわりに、掛けられた式 \(g(x)\) を別に持ち、それに \((1+a_i x)\) を掛けることを考える。このとき、\([x^i] f(x)g(x)\) を求めることになるが、これは \(O(\text{deg}(g))\) で求められる。ただし、\(\text{deg}(g)\) は \(g\) の次数とする。

更新は \(O(N)\) 回行われるので、\(g(x)\) が \(O(N)\) 次式になってしまう。そこで、\(\text{deg}(g) = \Theta(B)\) となったときに、\(f(x) \leftarrow f(x)g(x), g(x) \leftarrow 1\) と更新することで \(g\) の次数を下げる。

\(f(x) \leftarrow f(x)g(x)\) の操作は \(O\left(\frac{N}{B}\right)\) 回発生し、各操作に \(O(N\log N)\) 時間かかるので、合計で \(O\left(\frac{N^2 \log N}{B}\right)\) 時間かかる。また、\(g(x)\) に \((1 + a_i x)\) を掛けるとき、\(\text{deg}(g) = O(B)\) であるから、合計で \(O(NB)\) 時間かかる。

\(B = \Theta\left(\sqrt{N \log N}\right)\) とおくと、時間計算量 \(O \left(N \sqrt{N \log N}\right)\) となる。

実装例: https://atcoder.jp/contests/abc281/submissions/37186590