[NOIP2015 普及組] 求和

題目背景

NOIP2015 普及組 T3

題目描述

一條狹長的紙帶被均勻劃分出了\(n\)個格子，格子編號從\(1\)到\(n\)。每個格子上都染了一種顏色\(color_i\)用\([1,m]\)當中的一個整數表示），並且寫了一個數字\(number_i\)。

定義一種特殊的三元組：\((x,y,z)\)，其中\(x,y,z\)都代表紙帶上格子的編號，這裡的三元組要求滿足以下兩個條件：

1. \(xyz\)是整數,\(x<y<z,y-x=z-y\)

2. \(colorx=colorz\)

滿足上述條件的三元組的分數規定為\((x+z) \times (number_x+number_z)\)

輸入格式

第一行是用一個空格隔開的兩個正整數\(n\)和\(m,n\)表紙帶上格子的個數，\(m\)表紙帶上顏色的種類數。

第二行有\(n\)用空格隔開的正整數，第\(i\)數字\(number\)表紙帶上編號為\(i\)格子上面寫的數字。

第三行有\(n\)用空格隔開的正整數，第\(i\)數字\(color\)表紙帶上編號為\(i\)格子染的顏色。

輸出格式

一個整數，表示所求的紙帶分數除以\(10007\)所得的餘數。

樣例 #1

樣例輸入 #1

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

樣例輸出 #1

82

樣例 #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

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 $的顏色；

對 於 全 部 \(10\) 組 數 據 ， \(1 ≤ n ≤ 100000, 1 ≤ m ≤ 100000, 1 ≤ color_i ≤ m,1≤number_i≤100000\)

思路

一開始想到\(O(n^2)\)的演算法。

\(\because y-x=z-\)
\(\therefore x + z=2\times y\)

所以暴力列舉\(x,z\)就好了。但是肯定會\(TLE\)。

我們把所有數按顏色分成\(m\)組，然後為了列舉下標再開兩個位置來判斷下標的奇偶性（因為前面\(x+z=2\times y\)，所以\(x,z\)奇偶性相同）。

假設一組裡的數分別是\(x_1,x_2,\cdots,x_k\)，下標是\(y_1,y_2,\cdots,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 + \cdots + y_1 + y_k) + x_2\times(y_2 + y_1 + y_2 + y _ 3 + \cdots + y_2 + y_k) + \cdots + x_k\times(y_k + y_1 + y_k + y _ 2 + \cdots + y_k + y_{k-1})\)
\(~~~~~~~~~~~~~=x_1\times(y_1 \times (k - 2) + \sum\limits_{i=1}^k{y_i}) + x_2\times(y_2 \times (k - 2) + \sum\limits_{i=1}^k{y_i}) + \cdots + x_k\times(y_k \times (k - 2) + \sum\limits_{i=1}^k{y_i})\)

這裡每一個式子裡都有\(\sum\limits_{i=1}^k{y_i})\)，所以我們可以提前與處理一下，加快速度。

我們可以列舉所有數，第\(i\)數都加上\(x_i\times(y_i \times (k - 2) + \sum\limits_{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;
}