【集训队作业】COUNTARI

题目描述

给定一个长度为$n$的序列$\{a_n\}$
问有多少个三元组$(i, j, k)$满足，$a_i, a_j, a_k$是等差数列。

$n\leq 10^5, a_i\leq 3\times 10^4$

---

分析

考虑分块。
假如三个数都在当前块中，直接枚举并查找即可，$O(\frac{n}{B}B^2)=O(nB)$
假如只有两个数在当前块中，则分第三个数在左还是右，同样查找一下即可。$O(nB)$
假如有一个数在当前块中，则问题转化为从左边选一个右边选一个加起来等于枚举的$a_i$，FFT即可。$O(\frac{n}{B} m\log m)$

取$B=\sqrt{n}$就可以达到$O( n\sqrt{n} + \sqrt{n}m\log m$的时间复杂度。