題目的難度顏色使用 Luogu 上的分級，由簡單到困難分別為 🔴🟠🟡🟢🔵🟣⚫。

🔗 P2671 [NOIP 2015 普及组] 求和

Problem Statement

題目簡述

給定長度為 $n$ 的序列：每個位置 $i$ 有數值 $a_i$ 與顏色 $c_i$ 。

考慮所有位置的三元組 $(x, y, z)$ （ $1 \le x < y < z \le n$ ）且滿足 $y-x=z-y$ （等差），並且 $c_x=c_z$ 。

對每個符合條件的三元組，貢獻為 $(x+z)\cdot(a_x+a_z)$ ，求所有貢獻總和對 $10007$ 取模的結果。

Constraints

約束條件

$1 \le n \le 100000$
$1 \le a_i \le 10000$
$1 \le c_i \le m$ （ $m$ 為顏色種類數）
取模： $10007$

思路：貢獻法 + 分組前綴和

題目看起來在枚舉三元組 $(x,y,z)$ ，但真正出現在公式中的只有 $x,z$ ；中間點 $y$ 只用來限制「 $x,z$ 之間必須能形成等差」。

1) 先把條件化成「可配對」

由 $y-x=z-y$ 可得：

2y=x+z

因為 $y$ 必須是整數位置，所以 $x+z$ 必須是偶數，等價於：

關鍵化簡

$x$ 與 $z$ 的奇偶性相同（同為奇數或同為偶數）。

另外題目要求 $c_x=c_z$ ，因此任何有效三元組都對應到一個「同顏色 + 同奇偶」的配對 $(x,z)$ （且 $x<z$ ），而中間點 $y$ 由 $y=(x+z)/2$ 唯一決定。

於是可以把所有位置依照「顏色」與「位置奇偶」分組；同一組內任取兩個不同位置 $(x,z)$ 都能形成一個合法三元組（配對一次）。

2) 把每一組的總貢獻拆開計算

對於同一組中的一對 $(x,z)$ （ $x<z$ ），貢獻是：

(x+z)(a_x+a_z)

展開後：

x a_x + z a_z + x a_z + z a_x

可以把它理解成兩種部分：

每個元素的自項： $x a_x$ 與 $z a_z$
每個位置會和同組內其它 $m-1$ 個位置配對，因此自項總和就是

\sum_{t\in\text{group}} (t\cdot a_t)\cdot (m-1)

交叉項： $x a_z + z a_x$
這部分可以用「枚舉右端點、維護左端點前綴」：

當掃描到目前位置 $z$ 時，所有在它左邊的 $x$ 都會貢獻

\sum_{x<z} (x\cdot a_z + z\cdot a_x)= a_z\cdot\sum_{x<z}x + z\cdot\sum_{x<z}a_x

所以只要在每組維護兩個前綴和：

需要維護的量

左側位置編號總和 $\sum x$
左側數值總和 $\sum a_x$

每加入一個新元素，就能在 $\mathcal{O}(1)$ 更新交叉項並累加答案。

複雜度分析

時間複雜度： $\mathcal{O}(n)$
空間複雜度： $\mathcal{O}(n)$

Code

from collections import defaultdict

MOD = 10007


def solve():
    n, m = map(int, input().split())

    A = list(map(int, input().split()))
    C = list(map(int, input().split()))
    assert len(A) == len(C) == n

    groups = defaultdict(list)
    for i, (x, c) in enumerate(zip(A, C)):
        groups[(c, i & 1)].append((i + 1, x))

    ans = 0
    for g in groups.values():
        m = len(g)
        s1 = s2 = 0
        for a, b in g:
            ans += a * b * (m - 1)  # 每個元素本身的貢獻
            ans += a * s2 + b * s1  # 枚舉右維護左
            ans %= MOD
            s1 += a
            s2 += b
    print(ans)


if __name__ == "__main__":
    solve()