85pts求助
查看原帖
85pts求助
590925
_x_y_楼主2022/10/26 19:46
#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
#define int long long
using namespace std;
const int mod = 998244353;
const int N = 1e6 + 10;
int n, kk, a[N];
long long ans;
long long ansz, ansf;
int k;
void solvei1(){//m = 1
	for(int i = 1; i <= n; i++)
		ans = (ans + 1ll * (a[i] + 1) * (a[i] + 1) % mod) % mod;
}
void solvei2(){//m = 2
	for(int i = 1; i < k; i++)
		ans = (ans + 1ll * (a[i] + 1) * (a[i] + 1) % mod) % mod;
	for(int i = k; i <= n; i++)
		ans = (ans + 1ll * (a[i] + 2) * (a[i] + 2) % mod) % mod;
}
long long sum[N];
signed main(){
	cin >> n >> kk;
	for(int i = 1; i <= n; i++){
		cin >> a[i];
		sum[i] = (sum[i-1] + a[i]) % mod;
	}
	k = n + 1;
	for(int i = 1; i <= n; i++)
		if(a[i] >= 0){
			k = i;
			break;
		}
	for(int i = 1; i < k; i++)
		ansf = (ansf + 1ll * (a[i] + 1) * (a[i] + 1) % mod) % mod;
	for(int i = k; i <= n; i++)
		ansz = (ansz + 1ll * a[i] * a[i] % mod) % mod;
	solvei1();
	if(kk == 1){
		cout << ans << endl;
		return 0;
	}
	solvei2();
	//k-第一个"正数"
	for(int i = 3; i <= kk; i++){//i -> m
		while(k > 1 && a[k-1] + i >= abs(a[k-1])){
			k--;
			ansz = (ansz + 1ll * a[k] * a[k] % mod) % mod;
			ansf = (ansf - 1ll * (a[k] + 1) * (a[k] + 1) % mod) % mod;
		}
		int tansz = (ansz + 2ll * i * (sum[n] - sum[k - 1]) % mod + 1ll * i * i % mod * (n - k + 1) % mod) % mod;
		ans = (((ans + ansf) % mod) + tansz) % mod;
	}
	cout << ans << endl;
	return 0;
}
2022/10/26 19:46
加载中...