[ABC399-F] Range Power Sum 题解

题意就不多说了。

很明显，能打出如下暴力：

for (int i = 1; i <= n; ++i)
	for (int j = 0; j <= i; ++j)
		(ans += ppow(pre[i] - pre[j], k)) %= mod;

但在 \(n \le 2 \times 10^5\) 的情况下无法通过此题。

我们细细观察 \((pre_i - pre_j)^k\)。

谜团解开了。

形如 \((a-b)^k\) 的式子，可拆分为 \({\sum}_{m=0}^k\left(-1\right)^m\times\binom{k}{m}\times a^{k-m} \times b^m\)

代码转变为：

for (int i = 1; i <= n; ++i) {
	for (int j = 0; j <= i; ++j) {
		ll ret = 0;
		for (int o = 0; o <= k; ++o) {
			ret += ppow(-1, o) * C(k, o) * ppow(pre[i], k - o) * ppow(pre[j], o);
		ans += ret; ans %= mod;
    }
}

交换循环 \(i,j\) 的位置。

for (int j = 0; j <= i; ++j) {
    for (int i = 1; i <= n; ++i) {
		ll ret = 0;
		for (int o = 0; o <= k; ++o) {
			ret += ppow(-1, o) * C(k, o) * ppow(pre[i], k - o) * ppow(pre[j], o);
		ans += ret; ans %= mod;
    }
} 

每次循环 \(ret\) 都会加上 \(-1^o \times \binom{k}{o} \times pre_i^{k - o} \times pre_j^o\)。

而前半部分是定值，只有 \(pre_j^o\) 会改变。

做个前缀和，问题迎刃而解了。

时间复杂度为 \(\mathcal{O}(n \ k \ log \ k)\)

Code：

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair <int, int>;
const int N = 1e6 + 10;
const ll mod = 998244353;
ll a[N], pre[N], sum[N][11];
ll f[N], inv[N];
inline ll ppow(ll a, ll b) {
	ll ret = 1;
	while (b) {
		if (b & 1) (ret *= a) %= mod;
		(a *= a) %= mod; b >>= 1;
	}
	return ret;
}
inline ll C(int n, int k) {
	return f[n] * inv[k] % mod * inv[n - k] % mod;
}
int main() {
	ios::sync_with_stdio(0);
	cin.tie(0); cout.tie(0);
	ll n, k; cin >> n >> k;
	inv[0] = f[0] = 1;
	for (ll i = 1; i <= 10; ++i) {
		f[i] = f[i - 1] * (ll)i % mod;
		inv[i] = ppow(f[i], mod - 2);
	}
	for (int i = 1; i <= n; ++i) {
		cin >> a[i];
		pre[i] = (pre[i - 1] + a[i]) % mod;
		for (int j = 0; j <= k; ++j)
			sum[i][j] = (sum[i - 1][j] + ppow(pre[i], j)) % mod;
	}
	ll ans = 0;
	for (int i = 1; i <= n; ++i) {
		for (int j = 0; j <= k; ++j) {
			ll ret = C(k, j) % mod * ppow(pre[i], k - j) % mod * ((sum[i][j] + ppow(pre[0], j)) % mod) % mod;
			if (j % 2 == 0) {
				(ans += ret) %= mod;
			} else {
				ans = (ans - ret + mod) % mod;
			}
		}
	}
	cout << ans;
	return 0;
}