ABC399F - Range Power Sum

問題

リンク 長さ$N$の整数列$A$と$K$が与えられるので、以下の式ををmod 998で求めて下さい。

$$\sum_{1\leq l\leq r\leq N}(\sum_{_l\leq i \leq r} A_i)^K$$

解法

とりあえず与えられた式を理解しやすいように書き換えると、

$$\sum_{l=1}^N\sum_{r=l}^N(\sum_{i=l}^r A_i)^K$$

という３重シグマの形になる。が、一番内側のシグマについては区間和を表しているので、累積和を用いることで容易に外せる。数列$S$を$S_0 = 0$, $S_i = A_1 + \cdots + A_i$とすると、

$$\sum_{l=1}^N\sum_{r=l}^N(S_r - S_{l-1})^K$$

となる。どう考えても$K$乗の部分が厄介なので、二項展開¹を用いてこれを外すと、以下のようになる。

$$\sum_{l=1}^N\sum_{r=l}^N\sum_{p=1}^K \binom{K}{p}{S_r}^{K-p}(-S_{l-1})^p$$

3重シグマとなったが、これらを交換しても答えは変わらない。また、シグマから外側に出せるものは出してしまうことにする。これらを適用すると、式は以下のように変形できる。

$$\sum_{p=1}^K \binom{K}{p} \sum_{l=1}^N (-S_{l-1})^p \sum_{r=l}^N {S_r}^{K-p}$$

ここで、$\sum_{r=l}^N {S_r}^{K-p}$は${S_i}^x$の累積和を前計算しておくことによって、各$l$に対して$O(1)$で計算ができる。よって全体で$O(NK)$で解くことが可能で、この問題の制約下では十分高速になる。

提出

提出(C++23, 45ms)

using mint = modint998244353;
mint ncr(int n, int r) {
    mint res = 1;
    rep(i,r) {
        res *= (n-i);
    }
    rep(i,r) {
        res /= (i+1);
    }
    return res;
}

void solve() {
    int n, k; cin >> n >> k;
    auto a = i64_vec_IN(n);
    vector<mint> csum(n+1, 0);
    rep(i,n) csum[i+1] = csum[i] + a[i];


    vector<vector<mint>> ccsum(k+1, vector<mint>(n+2, 0));
    rep(p, k+1) {
        rep(i, n+1) {
            ccsum[p][i+1] = ccsum[p][i] + csum[i].pow(p);
        }
    }

    mint ans = 0;
    rep(p,k+1){
        mint s = 0;
        rep(l,n) {
            s += (csum[l] * (-1)).pow(p) * (ccsum[k-p][n+1] - ccsum[k-p][l+1]);
        }
        ans += s * ncr(k,p);
    }
    cout << ans.val() << '
';
}

Footnotes

1. https://rikeilabo.com/commentary-binomial-theorem ↩