P4062 [Code+#1]Yazid 的新生舞会 树状数组维护三阶前缀和

P4062 [Code+#1]Yazid 的新生舞会 树状数组维护三阶前缀和

神仙树状数组\线段树题

P4062 [Code+#1]Yazid 的新生舞会

发现偏序性质

我们希望统计“新生舞会”子区间，先要发现它的性质。

从简单的情况入手，例如考虑只有 0 和 1 的序列，如果统计 1 的“新生舞会”子区间，那么条件是 $cnt(l,r)>\lfloor \frac{r-l+1}{2}\rfloor$ . 自然地，我们想到维护一个前缀和快速求得 $cnt(l,r)$ ，记 $S_i$ 表示 $[1,i]$ 中 1 的个数，那么上式改写为：
$$\begin{array}{rl}{S}_r-S_{l-1}& >\frac{r-l+1}{2}\\ 2S_r-r& >2S_{l-1}-(l-1)\end{array}$$
于是我们发现，只要统计满足以上偏序关系的偏序对的个数就能得到 1 的“新生舞会”子区间个数。

利用等差序列性质改进算法

但是如果对每一个数都做一遍统计，不难发现其复杂度是 $O(n^2)$ 的，我们需要改进复杂度。

继续从简单情况入手，我们观察 $2S_i-i$ 的性质：

$i$	1	2	3	4	5	6	7	8
$a_i$	0	0	1	0	1	1	0	0
$2S_i-i$	-1	-2	-1	-2	-1	0	-1	-2

我们可以发现一个重要性质：$2S_i-i$ 总是以 $a_i=1$ 的位置为开头或者结尾组成一段公差为 $-1$ 的等差序列。

可以发现 $[1,2]$，$[3,4]$，$[5,5]$，$[6,8]$ 的 $2S_i-i$ 都组成了等差序列，那么如何利用这个性质？首先一个显然的结论是等差序列内部不会对答案产生贡献，因为内部总是递减，一定不满足之前推出的偏序关系，所以我们如果考虑插入等差序列（每一个数都作为一段等差序列的开头或者结尾 $O(n)$） 并查询等差序列与前面的等差序列产生的贡献（期望 $O(n\log n)$），就可能解决此题。

思考插入等差序列和查询答案的实现细节

现在做一个小推广，假设统计以下序列中，5 的“新生舞会”子区间

$i$	1	2	3	4	5	6	7	8
$a_i$	1	2	5	3	5	5	4	2

显然将 5 看作 1，其它数看作 0 就是上一个表格的情况。插入一段由 $a_i$ 中 $[3,4]$ 对应的 $2S_i-i$ 组成的等差序列 等价于给 $2S_i-i$ 的值域区间 $[-1,-2]$ 加上权值 1 ，也就是区间加。查询等差序列 $[3,4]$ 与前面产生的贡献就是：求等差序列中，每一个数与前面满足偏序关系的数的个数之和。例如这个例子，就是统计前面与 $-1$ 满足偏序关系的 $2S_i-i$ 的个数加上前面与 $-2$ 满足偏序关系的 $2S_i-i$ 的个数。

那么用简洁的表达式表达，就是：

记 $v(x)$ 为 $2S_i-i=x$ 的 $i$ 个数，也就是 $x$ 的权值；令 $p=2S_l-l$ 为等差序列开头的那个数 ,令 $q=2S_r-r$ 为等差序列结尾的那个数，那么由 $[l,r]$ 中的 $ 2S_i-i$ 组成的等差序列与前面的等差序列产生的贡献就是：
$$\sum _{q⩽k⩽p}\sum _{t<k}v(t)$$
为了提高运算速度，考虑用前缀和优化，并用树状数组维护前缀和。记 $g(x)={\sum }_{k⩽x}v(k)$ ，是 $2S_i-i⩽x$ 的 $i$ 的个数，也就是 $v(x)$ 的前缀和；记 $h(x)={\sum }_{k⩽x}g(k)$ ，表示对于所有 $2S_i-i⩽x$ 的 $2S_i-i$ ，满足 $2S_j-j⩽2S_i-i$ 的 $j$ 的个数之和，也就是 $g(x)$ 的前缀和；用上面的例子，我们知道查询 $[3,4]$ 组成的等差序列与前面的贡献就是 $h((-1)-1)-h((-2)-2)$ 。普遍意义下，查询开头为 $p=2S_l-l$ ，结尾为 $q=2S_r-r$ 的等差序列与前面的贡献是 $h(p-1)-h(q-2)$ ，这样就能求得答案。

于是我们可以制作下表：

$x=S_i-i$	-3	-2	-1	0	1
$v(x)$	0	1	1	0	0
$g(x)$	0	1	2	2	2
$h(x)$	0	1	3	5	7

然后考虑加入一段等差序列的问题：上面也已经说了，插入一段以 $[l,r]$ 中的 $ 2S_i-i$ 组成的等差序列等价于在 $2S_i-i$ 的权值域 $v$ 中对区间 $[2S_r-r,2S_l-l]$ 区间加 1.于是我们考虑对 $v$ 区间加 1 后，对 $h$ 有什么影响就好了，这里的思考思路参考《挑战程序设计竞赛》中处理树状数组区间加和区间查询的思路。

树状数组维护三阶前缀和

现在我们将目光转向 $v,g,h$ 三个函数，以下的 $[l,r]$ 都是指三个函数的定义域 $x=S_i-i$ 上的区间（前面的是指 $a_i$）

考虑令满足 $l⩽x⩽r$ 的 $v(x)$ 全部加 1.

记操作前的 $x$ 对应的 $g$ 函数为 $g(x)$ ，操作后为 $g^{\prime }(x)$ ，$h$ 函数同理。分类讨论：

• $x<l$ ，则 $g(x)$ 和 $h(x)$ 不受影响，即 $g^{\prime }(x)=g(x),h^{\prime }(x)=h(x)$

• $l⩽x⩽r$ ，则
  $$\begin{array}{rl}{g}^{\prime }(x)& =g(x)+x-l+1\\ h^{\prime }(x)& =h(x)+\sum _{k=l}^x(k-l+1)\\ & =h(x)+\frac{(x-l+1)(x-l+2)}{2}\\ & =h(x)+\frac{1}{2}\left(x^2+(3-2l)x+l^2-3l+2\right)\end{array}$$
  可以借用上面给出的表格与例子推一推。

• $r<x$ ，则
  $$\begin{array}{rl}{g}^{\prime }(x)& =g(x)+r-l+1\\ h^{\prime }(x)& =h(x)+\sum _{k=l}^r(k-l+1)+(x-r)(r-l+1)\\ & =h(x)+\frac{(r-l+2)(r-l+1)}{2}+(x-r)(r-l+1)\\ & =h(x)+\frac{1}{2}\left(2(r-l+1)x+(r-l+2)(r-l+1)-2r(r-l+1)\right)\end{array}$$

然后我们发现，变化量与 $x^2$ ，$x$ 和常数项（$x^0$）有关，于是用三个树状数组 $\mathrm{b}\mathrm{i}\mathrm{t}0$，$\mathrm{b}\mathrm{i}\mathrm{t}1$，$\mathrm{b}\mathrm{i}\mathrm{t}2$ 分别维护这三项，具体做法是：

对于 $[l,r]$ 区间加 1：

• 令 $\mathrm{b}\mathrm{i}\mathrm{t}0(l)$ 加上 $l^2-3l+2$；令 $\mathrm{b}\mathrm{i}\mathrm{t}1(l)$ 加上 $3-2l$ ；令 $\mathrm{b}\mathrm{i}\mathrm{t}2(l)$ 加上 $1$
• 令 $\mathrm{b}\mathrm{i}\mathrm{t}0(r+1)$ 加上 $(r-l+2)(r-l+1)-2r(r-l+1)-(l^2-3l+2)$ ；令 $\mathrm{b}\mathrm{i}\mathrm{t}1(\mathrm{r}+1)$ 加上 $2(r-l+1)-(3-2l)$ ；令 $\mathrm{b}\mathrm{i}\mathrm{t}2(r+1)$ 加上 $-1$

查询的时候，例如查 $h(x)$ 就是
$$h(x)=\frac{\mathrm{Query}(\mathrm{b}\mathrm{i}\mathrm{t}0,x)+\mathrm{Query}(\mathrm{b}\mathrm{i}\mathrm{t}1,x)\times x+\mathrm{Query}(\mathrm{b}\mathrm{i}\mathrm{t}2,x)\times x^2}{2}$$
这里乘上 $\frac{1}{2}$ 是因为我们刚刚修改的时候没有乘上 $\frac{1}{2}$ .

以上式子还要慢慢理解。这里再放出一张图，用于理解二阶前缀和的变化，也就是 $g(x)$ 的变化，至于 $h(x)$ 的变化，确实需要一定的数学功底与思维能力理解（这几条式子我也推了很久，推了 2h ，而且有几次推错了，甚至有几次多项式相乘展开都算错了，我数学太菜了…）

整理思路，解决问题

综上所述，我们就可以将记录一种数出现的所有位置，然后将原序列简化为仅包含 0 和 1 的序列，维护一下 $h(x)$ 并统计答案，逐个击破即可。最后要注意的就是不要傻傻地将负数作为下标了，设置一个偏移量 $\Delta$ 就可以避免数组越界地情况。时间复杂度 $O(n\log n)$

参考代码

#include<iostream>
#include<cstdio>
#include<vector> 

using namespace std;
typedef long long ll;
const ll dlt = 500003;
ll ans,t[5][1000010];
int n,type,x;
vector<ll > a[500005];

int lowbit(int num)
{
	return num&(-num);
}

void bitadd(int ty,int i,ll val)
{
	for(;i <= n + dlt;i += lowbit(i)) t[ty][i] += val;
	return ;
}

ll bitq(int ty,int i)
{
	ll res = 0;
	for(;i;i -= lowbit(i)) res += t[ty][i];
	return res;
}

ll BIT_Query(ll num)
{
	ll res = 0;
	res = ((bitq(0,num) + bitq(1,num)*num + bitq(2,num)*num*num)/2);
	return res;
}

void BIT_Add(ll l,ll r,ll val)//维护三个树状数组
{
	ll tmp0 = l*l - (ll)3*l + (ll)2;
	ll tmp1 = (ll)3 - l - l;
	ll tmp2 = (ll)1;
	bitadd(0,l,tmp0*val);
	bitadd(1,l,tmp1*val);
	bitadd(2,l,tmp2*val);
	tmp0 = -tmp0 + (r - l + (ll)1)*(r - l + (ll)2) - (ll)2*r*(r - l + (ll)1) ;
	tmp1 = -tmp1 + (ll)2*(r - l + (ll)1);
	tmp2 = -tmp2;
	bitadd(0,r + 1,tmp0*val);
	bitadd(1,r + 1,tmp1*val);
	bitadd(2,r + 1,tmp2*val);
	return ;
}

void solve(ll num)
{
	ll cnt = 0;ll lt = 0;ll rt = 0;
	BIT_Add(dlt,dlt,1);\\dlt 是偏移量
	if(a[num][0] > 1) lt = -a[num][0] + 1 ,rt = -1,BIT_Add(lt + dlt,rt + dlt,1); 
	for(int i = 0;i < a[num].size() - 1;i ++)//插入，查询，统计答案
	{
		cnt ++;rt = cnt + cnt - a[num][i];lt = cnt + cnt - a[num][i + 1] + 1;
		ans += BIT_Query(rt - 1 + dlt) - BIT_Query(lt - 2 + dlt);
		BIT_Add(lt + dlt,rt + dlt,1);
	}
	cnt = 0;
	BIT_Add(dlt,dlt,-1);
	if(a[num][0] > 1) lt = - a[num][0] + 1,rt = -1,BIT_Add(lt + dlt,rt + dlt,-1);
	for(int i = 0;i < a[num].size() - 1;i ++)//清空树状数组
	{
		cnt ++;lt = cnt + cnt - a[num][i + 1] + 1;rt = cnt + cnt - a[num][i];
		BIT_Add(lt + dlt,rt + dlt,-1);
	}
	return ;
}

int main()
{
	scanf("%d%d",&n,&type);
	for(int i = 1;i <= n;i ++)
	{
		scanf("%d",&x);
		a[x].push_back(i);//记录位置
	}
	for(int i = 0;i <= n;i ++)//逐个击破
		if(!a[i].empty())
		{
			a[i].push_back(n + 1);
			solve(i);
		}
	printf("%lld",ans);
	return 0;
 }