[ACNOI2022]数据结构之痛

题目

题目背景
《扬州慢·为卷爷吊打有感》

西师附中，联考之盟，解鞍静坐如松。过三时有半，得分数尚低。至揭榜大势已去，但见卷爷，独占榜一。渐黄昏，速补题毕，势惊游龙。

泥人有情，连日见虐起恨心。纵抱团相依，气血耗尽，卷爷亦赢。机房容颜未改，独缺予、无人送行。遇当年卷爷，已是今日国集！

题目描述
记 $f(b_1,b_2,\dots ,b_m)={\sum }_{i=1}^{m-1}\omega \left({\sum }_{j=1}^i{b}_j\right)$，其中 ω ( x ) = min ⁡ { x ,    k − x }    ( 0 ⩽ x < k ) \omega(x)=\min\{x,\;k{-}x\}\;(0\leqslant x<k) ω(x)=min{x,k−x}(0⩽x<k) 且 $\omega (x)=\omega (x\mathrm{mod}k)$ 。但是，若 ${\sum }_{j=1}^m{b}_j$ 不是 $k$ 的倍数，则定义 $f=-1$ 。

现在给出序列 { a } \{a\} {a}，求 ${\sum }_{i=1}^n{\sum }_{j=i}^{n}f(a_i,a_{i+1},\dots ,a_j)$，答案对 $998244353$ 取模。

数据范围与约定
$n⩽10^6$ 而 $k⩽10^8$ 。

思路

也许我的思路过分局限了。我首先就考虑（猫树）分治。不难发现当区间 $[l,r]$ 和为 $k$ 的倍数时，$\omega$ 值等于 $[l,m)$ 的前缀计算（包括整个区间）和 $[m,r]$ 的后缀计算（不包括整个区间）。因为总和为 $k$ 时，前缀的 $\omega$ 值等于其补集（一个后缀）的 $\omega$ 值。

但是，想要算出来 “前缀值” 需要 $\mathcal{O}(n\log n)$，总复杂度 $\mathcal{O}(n{\log }^{2}n)$，无法通过。

又考虑固定右端点。我本以为，多个左端点对应的 $\omega$ 的变化比较杂乱，难以搞定；我们永远滴神 $\mathsf{Rainybunny}$ 却告诉我，其实就是个 线性变换，可以用矩阵维护的。维护三元向量 $(w,r,1)$，其中 $w$ 对应当前 $\omega$ 而 $r$ 为区间和模 $k$ 的结果。那么，右端点移动，等价于将 $r$ 进行整体加，然后将 $r$ 超过 $k$ 的部分减去 $k$ 。可以用数据结构维护，比如 $\mathtt{treap}$ 。对于 $w$ 的变化，要么是 $+r$，要么是 $+k-r$，都可以写成矩阵。所以其实可以直接 $\mathcal{O}(n\log n)$ 。

不幸的是，数据结构题不能忽略常数。据称，平衡树的常数 $\gg$ 线段树 $>$ 树状数组。又乘了矩阵的大常数，所以 $\mathsf{Rainybunny}$ 没能通过本题，他能够因此长一智，我打心底里感到欣慰。

我还是太呆了。有用的 $f[l,r]$ 肯定满足前缀和 $s_r=s_{l-1}$ 。不难发现，若 $[l,m)$ 是合法的区间，$[m,r)$ 是合法区间，则 $[l,r)$ 合法且 $f[l,r)=f[l,m)+f[m,r)$ 。跟前面的分治做法用到相同的结论。

所以我们只需要对 “本原” 合法串计算 $f$ 值，乘上系数即可。只有 $\mathcal{O}(n)$ 个询问，可以离线后用芬威克树解决。记 $s_i$ 为前缀和数组，则 $f[l,r]={\sum }_{i=l}^r\omega (s_i-s_{l-1})$ 。只需求 ${\sum }_{i=1}^{\kappa }\omega (s_i-\lambda )$ 然后作差。用一个权值线段树（事实上是树状数组）记录前缀 { s i } \{s_i\} {si​} 的情况，查询就分段查 $s_i\in [\lambda ,\lambda +\lfloor \frac{k}{2}\rfloor ]$ 和剩余部分。时间复杂度 $\mathcal{O}(n\log n)$，常数不大。

代码

#include <cstdio> // JZM yydJUNK!!!
#include <iostream> // XJX yyds!!!
#include <algorithm> // XYX yydLONELY!!!
#include <cstring> // (the STRONG LONG for loneliness)
#include <cctype> // DDG yydDOG & ZXY yydBUS !!!
# define rep(i,a,b) for(int i=(a); i<=(b); ++i)
# define rep0(i,a,b) for(int i=(a); i!=(b); ++i)
# define drep(i,a,b) for(int i=(a); i>=(b); --i)
typedef long long llong;
inline int readint(){
	int a = 0, c = getchar(), f = 1;
	for(; !isdigit(c); c=getchar())
		if(c == '-') f = -f;
	for(; isdigit(c); c=getchar())
		a = (a<<3)+(a<<1)+(c^48);
	return a*f;
}

const int MOD = 998244353;
inline void modAddUp(int &x, const int &y){
	if((x += y) >= MOD) x -= MOD;
}

const int MAXN = 1000006;
int cnt[MAXN], sum[MAXN], cnt_all, sum_all, n, k, v[MAXN];
int query(int l, int r){
	if(l <= r){ // a complete range
		llong res = llong(v[l]+k)*cnt_all-sum_all; // rt+k-i
		for(int i=r; i; i&=(i-1)) // i-rt
			res = (res-llong(k+(v[l]<<1))*cnt[i]+(sum[i]<<1))%MOD;
		for(int i=l-1; i; i&=(i-1)) // rt-i
			res = (res-(sum[i]<<1)+llong(v[l]<<1)*cnt[i])%MOD;
		return int(res < 0 ? res+MOD : res);
	}
	else{ // two ranges
		llong res = sum_all-llong(cnt_all)*v[l]; // i-rt
		for(int i=l-1; i; i&=(i-1)) // rt-i
			res = (res-(sum[i]<<1)+llong(v[l]<<1)*cnt[i])%MOD;
		for(int i=r; i; i&=(i-1)) // i+k-rt
			res = (res+(sum[i]<<1)+llong(k-(v[l]<<1))*cnt[i])%MOD;
		return int(res < 0 ? res+MOD : res);
	}
}

int s[MAXN], lst[MAXN], nxt[MAXN];
int lcnt[MAXN], rcnt[MAXN], ans[MAXN];
int main(){
	n = readint(), k = readint();
	rep(i,1,n) if((s[i] = readint()+s[i-1]) >= k) s[i] -= k;
	memcpy(v+1,s,(n+1)<<2), std::sort(v+1,v+n+2);
	int *_end = std::unique(v+1,v+n+2);
	memset(lst+1,-1,(n+1)<<2); // clear
	for(int i=0,r; i<=n; ++i){
		s[i] = int(std::lower_bound(v+1,_end,s[i])-v);
		for(int j=s[i]; j<=n+1; j+=(j&-j)) // modify
			++ cnt[j], modAddUp(sum[j],v[s[i]]);
		++ cnt_all, modAddUp(sum_all,v[s[i]]);

		const int want = v[s[i]]+(k>>1);
		if(want < k || want-k < v[1])
			r = int(std::lower_bound(
				v+1,_end,want+1)-v-1);
		else r = int(std::lower_bound(
			v+1,_end,want-k+1)-v-1);
		ans[i] = MOD-query(s[i],r); // fast

		if(lst[s[i]] != -1){
			const int &pre = lst[s[i]];
			modAddUp(ans[pre],query(s[pre],r));
			nxt[pre] = i, lcnt[i] = lcnt[pre]+1;
		}
		else lcnt[i] = 1; // start
		lst[s[i]] = i; // new position
	}
	int xjx = 0; ///< the final answer
	drep(i,n,0) // everything
		if(!nxt[i]) rcnt[i] = 1;
		else{ // calculate
			rcnt[i] = rcnt[nxt[i]]+1;
			xjx = int((xjx+llong(rcnt[nxt[i]])
				*lcnt[i]%MOD*ans[i])%MOD);
		}
	memset(rcnt+1,0,(n+1)<<2); // reuse
	llong bad = llong(n+1)*n>>1;
	rep(i,0,n) ++ rcnt[s[i]]; // bucket sort
	rep(i,1,n+1) bad -= llong(rcnt[i]-1)*rcnt[i]>>1;
	xjx = int((xjx-bad%MOD+MOD)%MOD);
	printf("%d\n",xjx);
	return 0;
}