G - Colorful Candies 2 Editorial by MMNMM
Polynomial Taylor Shift を使った考え方\(K\) を固定したとき、その答えが \(\displaystyle\left.\sum_{i=1}^C\left\{\dbinom{N}{K}-\dbinom{N-n_i}{K}\right\}\middle/\dbinom{N}{K}\right.\) であることまでは公式解説と同じです。 これの \(K=1,2,\ldots,N\) における値を Polynomial Taylor Shift と呼ばれる概念を用いることで計算します。
Polynomial Taylor Shift とは
多項式 \(\displaystyle f(x)=\sum_{i=0}^Na_ix^i\) と定数 \(c\) について\(\displaystyle f(x+c)=\sum_{i=0}^Nb_ix^i\) のことを Polynomial Taylor Shift といい、列 \((a_i)_{i=0}^N\) と定数 \(c\) が与えられたときにこれを満たす列 \((b_i)_{i=0}^N\) を畳み込み \(1\) 回と \(O(N)\) 回の演算で求める方法があります。
具体的な方法は式変形によって確認することができます。 \(f(x+c)\) は
\[\begin{aligned} \displaystyle f(x+c)&=\sum_{i=0}^Na_i(x+c)^i\\ &=\sum_{i=0}^Na_i\sum_{j=0}^i\dbinom{i}{j}c^{i-j}x^j\\ &=\sum_{j=0}^Nx^j\sum_{i=j}^N\dfrac{i!}{j!(i-j)!}a_ic^{i-j}\\ &=\sum_{j=0}^N\dfrac{x^j}{j!}\sum_{i=j}^N(a_ii!)\dfrac{c^{i-j}}{(i-j)!} \end{aligned}\]
と変形できます。 よって、各 \(j=0,\ldots,N\) について \(\displaystyle b_j=\dfrac{1}{j!}\sum_{i=j}^N(a_ii!)\dfrac{c^{i-j}}{(i-j)!}\) となります。 ここで \((a_i)_{i=0}^N,(b_i)_{i=0}^N\) を逆順に見ることで、
\[\displaystyle b_{N-j}=\dfrac{1}{(N-j)!}\sum_{N-i=0}^{N-j}(a_{N-i}(N-i)!)\dfrac{c^{(N-j)-(N-i)}}{((N-j)-(N-i))!}\]
と畳み込みの形になっていることがわかります。
事前に \(1!,2!,\ldots,N!\) およびその逆元を計算しておくことで(これは \(O(N)\) 時間で可能です)、\(1\) 度の畳み込みで \((b_i)_{i=0}^N\) を求めることができます。
残りの部分
今回の問題でこれをどのように使うのか解説します。
\(K\) に対する答えは \(\displaystyle \left.\sum_{i=1}^C\left\{\dbinom{N}{K}-\dbinom{N-n_i}{K}\right\}\middle/\dbinom{N}{K}\right.\) でした。 これを変形すると \(\displaystyle C-\left(1\middle/\dbinom{N}{K}\right)\sum_{i=1}^C\dbinom{N-n_i}{K}\) となります。 適切な前計算のもとで \(\dbinom{N}{K}\) は \(1\) つの\(K\) あたり \(O(1)\) 時間で計算ができるので、 \(\displaystyle f_K=\sum_{i=1}^C\dbinom{N-n_i}{K}\) が高速に求まればよいです。
ここで \(\displaystyle f(x)=\sum_{i=0}^Nf_ix^i\) を考えると、 これは
\[\begin{aligned} \displaystyle f(x)&=\sum_{i=0}^Nf_ix^i\\ &=\sum_{i=0}^N\sum_{j=1}^C\dbinom{N-n_j}{i}x^i\\ &=\sum_{j=1}^C\sum_{i=0}^N\dbinom{N-n_j}{i}x^i=\sum_{j=1}^C(1+x)^{N-n_j}\\ \end{aligned}\]
と変形でき、\(\displaystyle g(x)=\sum_{j=1}^Cx^{N-n_j}\) に対して \(g(x+1)\) と等しいです。 \(g(x)\) の係数の列は \(O(N)\) 時間で計算ができるので、これに上記のアルゴリズムを使うことで \((f_i)_{i=1}^N\) を \(O(N\log N)\) 時間で求めることができました。
実装は以下のようになります(実際の提出)。 この実装では \(1!,2!,\ldots,N!\) およびその逆元の前計算を \(O(N+\log\mathrm{MOD})\) で行っているので、全体の計算量は \(O(N\log N+\log\mathrm{MOD})\) です。
#include<bits/stdc++.h>
#include<atcoder/modint>
#include<atcoder/convolution>
int main(){
using namespace std;
using modint = atcoder::static_modint<998244353>;
unsigned long N;
cin >> N;
// 階乗前計算
vector<modint> fact{modint::raw(1)};
for(unsigned long i{1}; i <= N; ++i)fact.emplace_back(fact.back() * modint::raw(i));
vector<modint> ifact{fact.back().inv()};
for(unsigned long i{N}; i; --i)ifact.emplace_back(ifact.back() * modint::raw(i));
reverse(begin(ifact), end(ifact));
// n_i の計算
map<unsigned long, unsigned long> cnt;
for(unsigned long i{}, c; i < N; ++i){
cin >> c;
++cnt[c];
}
// g(x) の計算
vector<modint> a(N + 1);
for(const auto& [_, c] : cnt)++a[N - c];
// Taylor Shift : a -> b
// a_i -> a_i i!
for(unsigned long i{}; i <= N; ++i)a[i] *= fact[i];
reverse(begin(a), end(a));
// b_i i! <- (a_i i!) * (1 / i!)
auto b{atcoder::convolution(a, ifact)};
b.resize(N + 1);
reverse(begin(b), end(b));
// b_i <- b_i i!
for(unsigned long i{}; i <= N; ++i)b[i] *= ifact[i];
const auto C{modint::raw(size(cnt))};
for(unsigned long i{1}; i <= N; ++i)cout << (C - b[i] * ifact[N] * fact[i] * fact[N - i]).val() << endl;
return 0;
}
posted:
last update: