P2671 [NOIP2015 普及组] 求和

[NOIP2015 普及组] 求和

题目背景

NOIP2015 普及组 T3

题目描述

一条狭长的纸带被均匀划分出了 $n$ 个格子，格子编号从 $1$ 到 $n$ 。每个格子上都染了一种颜色 $c o l o r_{i}$ 用 $[1, m]$ 当中的一个整数表示），并且写了一个数字 $n u m b e r_{i}$ 。

定义一种特殊的三元组： $(x, y, z)$ ，其中 $x, y, z$ 都代表纸带上格子的编号，这里的三元组要求满足以下两个条件：

$x y z$ 是整数, $x < y < z, y - x = z - y$
$c o l o r x = c o l o r z$

满足上述条件的三元组的分数规定为 $(x + z) \times (n u m b e r_{x} + n u m b e r_{z})$ 。整个纸带的分数规定为所有满足条件的三元组的分数的和。这个分数可能会很大，你只要输出整个纸带的分数除以 $10, 007$ 所得的余数即可。

输入格式

第一行是用一个空格隔开的两个正整数 $n$ 和 $m, n$ 表纸带上格子的个数， $m$ 表纸带上颜色的种类数。

第二行有 $n$ 用空格隔开的正整数，第 $i$ 数字 $n u m b e r$ 表纸带上编号为 $i$ 格子上面写的数字。

第三行有 $n$ 用空格隔开的正整数，第 $i$ 数字 $c o l o r$ 表纸带上编号为 $i$ 格子染的颜色。

输出格式

一个整数，表示所求的纸带分数除以 $10007$ 所得的余数。

样例 #1

样例输入 #1

6 2
5 5 3 2 2 2
2 2 1 1 2 1

样例输出 #1

样例 #2

样例输入 #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

提示

【输入输出样例 1 说明】

纸带如题目描述中的图所示。

所有满足条件的三元组为： $(1, 3, 5), (4, 5, 6)$ 。

所以纸带的分数为 $(1 + 5) \times (5 + 2) + (4 + 6) \times (2 + 2) = 42 + 40 = 82$ 。

对于第 $1$ 组至第 $2$ 组数据， $1 \leq n \leq 100, 1 \leq m \leq 5$ ；

对于第 $3$ 组至第 $4$ 组数据， $1 \leq n \leq 3000, 1 \leq m \leq 100$ ；

对于第 $5$ 组至第 $6$ 组数据， $1 \leq n \leq 100000, 1 \leq m \leq 100000$ ，且不存在出现次数超过 $20$ 的颜色；

对于全部 $10$ 组数据， $1 \leq n \leq 100000, 1 \leq m \leq 100000, 1 \leq c o l o r_{i} \leq m, 1 \leq n u m b e r_{i} \leq 100000$

思路

一开始想到 $O (n^{2})$ 的算法。

$∵ y - x = z -$
$∴ x + z = 2 \times y$

所以暴力枚举 $x, z$ 就好了。但是肯定会 $T L E$ 。

我们把所有数按颜色分成 $m$ 组，然后为了枚举下标再开两个位置来判断下标的奇偶性（因为前面 $x + z = 2 \times y$ ，所以 $x, z$ 奇偶性相同）。

假设一组里的数分别是 $x_{1}, x_{2}, \dots, x_{k}$ ，下标是 $y_{1}, y_{2}, \dots, y_{k}$
那么答案 $= (x_{1} + x_{2}) \times (y_{1} + y_{2}) + (x_{1} + x_{3}) \times (y_{1} + y_{3}) + \dots$
$= x_{1} \times (y_{1} + y_{2} + y_{1} + y_{3} + \dots + y_{1} + y_{k}) + x_{2} \times (y_{2} + y_{1} + y_{2} + y_{3} + \dots + y_{2} + y_{k}) + \dots + x_{k} \times (y_{k} + y_{1} + y_{k} + y_{2} + \dots + y_{k} + y_{k - 1})$
$= x_{1} \times (y_{1} \times (k - 2) + \sum_{i = 1}^{k} y_{i}) + x_{2} \times (y_{2} \times (k - 2) + \sum_{i = 1}^{k} y_{i}) + \dots + x_{k} \times (y_{k} \times (k - 2) + \sum_{i = 1}^{k} y_{i})$

这里每一个式子里都有 $\sum_{i = 1}^{k} y_{i})$ ，所以我们可以提前与处理一下，加快速度。

我们可以枚举所有数，第 $i$ 数都加上 $x_{i} \times (y_{i} \times (k - 2) + \sum_{i = 1}^{k} y_{i})$ 即可，最后全部加上模上 $10007$ 即可
可以依据代码来理解，我觉得挺有必要。

代码

#include <iostream>
using namespace std;
const int N = 100010,MOD = 10007;
int n,m;
int a[N],color[N];
int s[N][2],sum[N][2];
int main () {
	cin >> n >> m;
	for (int i = 1;i <= n;i++) cin >> a[i];
	for (int i = 1;i <= n;i++) {
		cin >> color[i];
		s[color[i]][i % 2]++;
		sum[color[i]][i % 2] = (sum[color[i]][i % 2] + i) % MOD;
	}
	int ans = 0;
	for (int i = 1;i <= n;i++) ans = (ans + a[i] * (i * (s[color[i]][i % 2] - 2) % MOD + sum[color[i]][i % 2]) % MOD) % MOD;
	cout << ans << endl;
	return 0;
}