T4 求做法

T3 70pts ，WA on 14~15,17~20,有几个点居然输出为负

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N=1e6+5,mod=998244353;
int a[N],suf[N];
int read()
{
	int f=0,x=0;
	char ch=getchar();
	while(!isdigit(ch))
	{
		if(ch=='-')f=1;
		ch=getchar();
	}
	while(isdigit(ch))x=(x<<1)+(x<<3)+(ch^48),ch=getchar();
	return f?-x:x;
}
signed main()
{
	int n,k,sum=0;
	n=read(),k=read();
	for(int i=1;i<=n;i++)a[i]=read()+1,(sum+=a[i]*a[i])%=mod;
	for(int i=n;i>=1;i--)suf[i]=(suf[i+1]+a[i])%mod;
	int p=n,ans=0;
	for(int i=0;i<k;i++)
	{
		while(p&&abs(a[p]+i)>abs(a[p]))p--;
		(ans+=sum+2*suf[p+1]*i+(n-p)*i*i)%=mod;
	}
	cout<<ans;
	return 0;
}