rt 代码如下

#include<bits/stdc++.h>
using namespace std;
const long long mod = 998244353;
const long long MAXN = 1e6 + 5;
long long n, k;
long long a[MAXN];
long long sum1[MAXN];
long long sum2[MAXN];
long long tot;
long long ADD(long long a, long long b)
{
	return ((a % mod) + (b % mod)) % mod;
}
long long SUB(long long a, long long b)
{
	return ((a % mod) - (b % mod) + mod) % mod;
}
long long MUL(long long a, long long b)
{
	return ((a % mod) * (b % mod)) % mod;
}
long long get_ans(long long i, long long m)
{
	long long ans = -1e18;
	long long temp1, temp2, temp3;
	temp3 = ADD(SUB(n, ADD(i, 1ll)), 1ll);
	temp1 = ADD(ADD(sum2[i], MUL(sum1[i], 2ll)), i);
	temp2 = ADD(ADD(SUB(sum2[n], sum2[i]), MUL(SUB(sum1[n], sum1[i]), MUL(2ll, m))), MUL(temp3, MUL(m, m)));
	ans = max(ans, ADD(temp1, temp2));
	return ans;
}
int main()
{
	cin >> n >> k;
	for(int i = 1; i <= n; i ++)
	{
		cin >> a[i];
	}
	for(int i = 1; i <= n; i ++)
	{
		sum1[i] = ADD(sum1[i - 1], a[i]);
		sum2[i] = ADD(sum2[i - 1], MUL(a[i], a[i]));
	}
	long long t;
	tot = ADD(tot, get_ans(0, 1));
	for(int i = 1; i <= n; i ++)
	{
		if(a[i] >= -1)
		{
			t = i - 1;
			break;
		}
	}
	for(long long i = 2; i <= k; i ++)
	{
		while(t >= 1 && abs(a[t] + i) > abs(a[t] + 1))
		{
			t --;
		}
		tot = ADD(tot, get_ans(t, i));
	}
	cout << ((tot % mod) + mod) % mod << "\n";
	return 0;
}

提交记录