#数学 #二項展開 #累積和
リンク
長さ
\(N\)
の整数列
\(A\)
と
\(K\)
が与えられるので、以下の式ををmod 998で求めて下さい。
とりあえず与えられた式を理解しやすいように書き換えると、
\[\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)\) で解くことが可能で、この問題の制約下では十分高速になる。
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() << '\n';
}