洛谷 P2671 求和
题目描述
一条狭长的纸带被均匀划分出了\(n\)个格子，格子编号从\(1\)到\(n\)。每个格子上都染了一种颜色\(color\_i\)用\([1,m]\)当中的一个整数表示，并且写了一个数字\(number\_i\)。
定义一种特殊的三元组：\((x,y,z)\)，其中\(x,y,z\)都代表纸带上格子的编号，这里的三元组要求满足以下两个条件：
\(x,y,z\)是整数,\(x<y<z,y-x=z-y\)
\(color_x\)=\(color_z\)
满足上述条件的三元组的分数规定为\((x+z)\times(number\_x+number\_z)\)。整个纸带的分数规定为所有满足条件的三元组的分数的和。这个分数可能会很大，你只要输出整个纸带的分数除以\(10,007\)所得的余数即可。
输入格式
第一行是用一个空格隔开的两个正整数\(n\)和\(m\)。\(n\)表示纸带上格子的个数，\(m\)表纸带上颜色的种类数。
第二行有\(n\)用空格隔开的正整数，第\(i\)数字\(number\)表纸带上编号为\(i\)格子上面写的数字。
第三行有\(n\)用空格隔开的正整数，第\(i\)数字\(color\)表纸带上编号为\(i\)格子染的颜色。
输出格式
一个整数，表示所求的纸带分数除以\(10007\)所得的余数。
输入输出样例
输入 #1
6 2
5 5 3 2 2 2
2 2 1 1 2 1
输出 #1
82
输入 #2
15 4
5 10 8 2 2 2 9 9 7 7 5 6 4 2 4
2 2 3 3 4 3 3 2 4 4 4 4 1 1 1
输出 #2
1388
说明/提示
【输入输出样例 1 说明】
纸带如题目描述中的图所示。
所有满足条件的三元组为： \((1, 3, 5), (4, 5, 6)\)。
所以纸带的分数为\((1 + 5) \times (5 + 2) + (4 + 6) \times (2 + 2) = 42 + 40 = 82\)。
对于第 \(1\) 组至第 \(2\) 组数据， \(1 ≤ n ≤ 100, 1 ≤ m ≤ 5\)；
对于第 \(3\) 组至第 \(4\) 组数据， \(1 ≤ n ≤ 3000, 1 ≤ m ≤ 100\)；
对于第 \(5\) 组至第 \(6\) 组数据， \(1 ≤ n ≤ 100000, 1 ≤ m ≤ 100000\)，且不存在出现次数超过$20$20的颜色；
对 于 全 部 \(10\) 组 数 据 ， \(1 ≤ n ≤ 100000, 1 ≤ m ≤ 100000, 1 ≤ color\_i ≤ m,1≤number\_i≤100000\)
思路
我离做出这个题只差一步化简式子然后优化了，但还是没有想到。
其实前面的思路比较好想，暴力当然就是\(O(n^2)\)两层循环枚举，如果两个数颜色和奇偶性都相同，那就加和（\(y-x=z-y\)其实就是\(x+z=2y\)，所以奇偶性相同）。但是这样只能得到\(40pts\)，要思考正解。
通过两层循环可以发现，我们其实不需要管中间那个点是什么，只要确定两边的点即可。而两边的点又满足颜色相同和奇偶性相同，我们就可以把所有的格子分成\(2m\)段，每一种颜色的奇数号和偶数号分开，就有这么多段。会发现，在段内任意选两个数，都是满足题意的，就加和。但如果分开后还是暴力循环相当于没有任何优化，我就是卡在了这里。但是我们可以将式子写出来，然后化简，寻找优化的方法。
以其中一段为例，比如有一段里面有\(k\)个格子。然后第\(i\)个格子的值是\(x_i\)，编号是\(y_i\)，那么对答案的贡献就是
(x1+x2)*(y1+y2)+(x1+x3)*(y1+y3)+...+(x1+xk)*(y1+yk)+
(x2+x3)*(y2+y3)+(x2+x4)*(y2+y4)+.........
然后我们随便取个\(k\)，然后用手化简两组样例，就会发现这样的规律：
原式 = x1*(y1*(k-1)+y2+y3+...+yk)+x2*(y2*(k-1)+y1+y3+...+yk)+......
上式 = x1*(y1*(k-2)+y1+y2+...+yk)+x2*(y2*(k-2)+y1+y2+...+yk)+......
这样的话，我们只需要预处理出前缀和，然后扫一遍一开始的序列即可，时间复杂度为\(O(n)\)。
代码
#include<iostream>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<queue>
using namespace std;
typedef long long int ll;
const int mod=10007;
int n,m,ans;
int tot[200005],num[100005],color[100005],sum[200005];
int main()
{
scanf("%d%d",&n,&m);
for(int i=1;i<=n;i++){
scanf("%d",&num[i]);
}
for(int i=1;i<=n;i++){
scanf("%d",&color[i]);
}
for(int i=1;i<=n;i++){
if(i%2==0){
tot[color[i]]++;//记录单数颜色那一段的个数
sum[color[i]]=(sum[color[i]]+i)%mod;//预处理前缀和
}
else{
tot[color[i]+m]++;//+m防止重复，记录偶数颜色个数
sum[color[i]+m]=(sum[color[i]+m]+i)%mod;//预处理前缀和
}
}
for(int i=1;i<=n;i++){
if(i%2==0){
ans=(ans+num[i]*(i*(tot[color[i]]-2)%mod+sum[color[i]]))%mod;//套用公式
}
else{
ans=(ans+num[i]*(i*(tot[color[i]+m]-2)%mod+sum[color[i]+m]))%mod;
}
}
printf("%d\n",ans);
return 0;
}
