\(H < W\) であるとき、時間計算量 \(O(H^{2} + H \log(W))\) で解けます。

想定解を高速化することで、この計算量を達成できます。

想定解を C++ で書いたものがこちらです。

想定解 C++

dpC が A の場合の数、 dpS が \(\sum A\) の総和に対応しています。

これを、Binomial を使わず、式変形をして、階乗の逆数のみを用いて書いたものがこちらです。

階乗の逆数 C++

一部抜粋すると、以下のようになります。

void solve(){
    dpC[0] = 1;
    rep(rp, 0, H){
        vector<mint> n_dpC(W + 1), n_dpS(W + 1);
        rep(i, 0, W + 1){
            rep(j, i, W + 1){
                n_dpC[j] += dpC[i] * invfact[j - i + 1] * (j + 1);
                n_dpS[j] += dpS[i] * invfact[j - i + 1] * (j + 1);
                n_dpS[j] += dpC[i] * i * invfact[j - i + 2] * (j + 1);
            }
        }
        swap(n_dpS, dpS);
        swap(n_dpC, dpC);
        rep(i, 0, W + 1) dpS[i] += dpC[i] * (W - i);
    }
    mint ans = 0;
    rep(i, 0, W + 1) ans += dpS[i] * invfact[W - i];
    rep(i, 1, W + 1) ans *= i;
    cout << ans.val() << "\n";
}

\(x\) の関数 \(f(x), g(x)\) を \(\displaystyle f(x) = \sum x^{k} dpC[k], g(x) = \sum x^{k} dpS[k]\) と定義したとき、dp を更新する操作は、\(f, g\) を微分する操作や、\(x\) の関数をかける操作に対応します。

たとえば、

rep(i, 0, W + 1){
    rep(j, i, W + 1){
        n_dpC[j] += dpC[i] * invfact[j - i + 1];
        n_dpS[j] += dpS[i] * invfact[j - i + 1];
        n_dpS[j] += dpC[i] * i * invfact[j - i + 2]
    }
}

は、

\(f_{new} = \dfrac{d\left(f(x)\left(e^{x} - 1\right)\right)}{dx}\)

\(g_{new} =\dfrac{d\left(g(x)\left(e^{x} - 1\right)\right)}{dx} + \dfrac{d\left(\left(\dfrac{df(x)}{dx}\right)\left(e^{x} - x - 1\right)\right)}{dx}\)

と表せます。

rep(i, 0, W + 1) dpS[i] += dpC[i] * (W - i);

は、

\(g_{new} = g + W\cdot f(x) - x\dfrac{df(x)}{dx}\)

と表せます。

\(e^{x}\) などをそのまま持つのは扱いにくので、以下の条件を満たす \(X\) の多項式 \(F(X), G_{0}(X), G_{1}(X)\) を定義します。

• \(F(e^{x}) = f(x)\)
• \(G_{0}(e^{x}) + xG_{1}(e^{x}) = g(x)\)

\(H\) 回のループの中では、\(F, G_{0}, G_{1}\) の次数は高々 \(H\) に収まります。

これを用いるとたとえば、

\(f_{new} = \dfrac{d\left(f(x)\left(e^{x} - 1\right)\right)}{dx}\)

は

\(\dfrac{d\left(f(x)\left(e^{x} - 1\right)\right)}{dx} = \dfrac{dX}{dx}\cdot\dfrac{d\left(F(X)\left(X - 1\right)\right)}{dX} = X \cdot (F(X) + F'(X)(X - 1))\)

と変形できることから

\(\displaystyle F_{new}(X) = \sum_{k} ((k + 1)X^{k + 1} - kX^{k})\cdot [X^{k}](F(X))\)

などと表せます。

これらの \(F, G_{0}, G_{1}\) の更新は \(O(H)\) で行えるので、最終的な \(F, G_{0}, G_{1}\) は時間計算量 \(O(H^{2})\) で求められます。

最後の答えを求めるパートである、

mint ans = 0;
rep(i, 0, W + 1) ans += dpS[i] * invfact[W - i];
rep(i, 1, W + 1) ans *= i;
cout << ans.val() << "\n";

は、

\(ans =W! [x^{W}] g(x)e^{x}\)

と表せるので、

\[\displaystyle ans = W![x^{W}] (G_{0}(e^{x}) + G_{1}(e^{x}))e^{x} = \sum_{k}\left((k + 1)^{W}[X^{k}]G_{0}(X) + (k + 1)^{W - 1}[X^{k}]G_{1}(X)\right)\]

となります。これは、時間計算量 \(O(H\log(W))\) で計算できます。

以上で全体の時間計算量は \(O(H^{2} + H\log(W))\) となります。

実装例 C++ 3 sec