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: