Colored Midpoints

解説1

値 $x$ が $A$ に現れる回数を $c_x$ とする。

適当な正整数 $B$ を定め、各 $x$ について、$c_x \le B$ と $c_x \gt B$ に分けて考える。

• $c_x \le B$

  $A_i=A_j=x$ を満たす組 $(i,j)$ をすべて列挙し、残りの条件を満たすか判定すれば良い。

  すべての $x$ について、この処理の計算量は合計で $O(NB)$ である。

• $c_x \gt B$

  多項式

  $$F_x(z)=\sum_{i:A_i=x}z^i$$

  を考える。

  $F_x(z)^2$ の $z^{2m}$ の係数は、$A_i=A_j=x$ かつ $i+j=2m$ を満たす順序付きの組 $(i,j)$ の個数である。

  このうち $i\ne j$ である組はそれぞれ $2$ 回数えられており、$i=j$ である組は $A_m=x$ のときに限り $1$ 回数えられている。

  よって、値 $x$ に対応する答えは

  $$\sum_{m:S_m=\texttt{1}} \frac{[z^{2m}]F_x(z)^2-[z^m]F_x(z)}{2}$$

  となる。

  $F_x(z)^2$ は畳み込みにより $O(N\log N)$ 時間で計算できる。

  $c_x\gt B$ を満たす $x$ は高々 $N/B$ 種類なので、この場合の計算量は合計で $O(N^2\log N/B)$ である。

以上より、計算量は $O(NB+N^2\log N/B)$ となる。

$B=\Theta(\sqrt{N\log N})$ とすることで、$O(N\sqrt{N\log N})$ でこの問題を解くことができる。

解説2

以下、$A,S$ を $0$-indexed とする。

$i+j$ が偶数であることは、$i,j$ の偶奇が一致することと同値である。

そこで、$A_i$ の値と $i$ の偶奇が一致するものをグループ化し、それぞれについて条件を満たす組を数える。

同じグループに属する $i\lt j$ を $i=2a+p$, $j=2b+p$（$p\in\{0,1\}$、$a\lt b$）と表すと、

$$\frac{i+j}{2}=a+b+p$$

である。

$T[m]=S_m$ となる bitset $T$ を用意する。ただし、0, 1 はそれぞれビットの $0$, $1$ として扱う。

各グループについて、$A_{2b+p}=x$ を満たす $b$ に対して $U[b]=1$ となる bitset $U$ を用意する。ここで、$x$ はそのグループの $A_i$ の値である。

$a$ を昇順に処理し、$A_{2a+p}=x$ を満たすたびに以下の操作を行う。

1. $U[a]=0$ とする。
2. $(U\mathbin{\&}(T\gg(a+p)))$ の立っているビットの個数を答えに加算する。

このとき、

$$U[b]=1 \iff a\lt b \land A_{2a+p}=A_{2b+p}$$

$$(T\gg(a+p))[b]=1 \iff S_{a+b+p}=\texttt{1}$$

であるため、これによって正しい答えを得られる。

以上より、$w$ をwordsize として $O(N^2/w)$ でこの問題を解くことができる。