Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 293 G問題 Triple Index

AtCoder ABC293 ABC 600 pts G問題

問題
提出解答
問題の概要
- 制約
解法

問題

提出解答

問題の概要

長さ $N$ の整数列 $A=(A_1, \dots, A_N)$ がある.

次の $Q$ 個の質問に答えよ.

質問 : 以下を満たす整数の組の個数を求めよ.
- $l_q \leq i \lt j \lt k \leq r_q$
- $A_i=A_j=A_k$

制約

$1 \leq N \leq 2 \times 10^5$
$1 \leq Q \leq 2 \times 10^5$
$1 \leq A_i \leq 2 \times 10^5$
$1 \leq l_q \leq r_q \leq N$

解法

この問題について, 以下のことがわかる.

$A$ は固定である.
求めるべき範囲 $[l_q, r_q ]$ は最初に与えられている.
$[l_q, r_q ]$ に対する結果が分かっているとき, $[l_q \pm 1, r_q ], [l_q, r_q \pm 1 ]$ における結果は高速に計算できる (方法は後述).

このようなとき, Mo アルゴリズムによって高速に計算できる.

よって, 次の $2$ つについてがわかれば良い.

1つの区間 $[l_q, r_q ]$ における結果の計算方法
$[l_q, r_q ]$ の結果を利用して, $[l_q \pm 1, r_q ], [l_q, r_q \pm 1 ]$ の結果を高速に計算する方法

$A_l, \dots, A_r$ において, 整数 $a$ が $\chi(a)$ 個あるとする. このとき, 条件を満たす $(i,j,k)$ の個数は

$\displaystyle \sum_{a} \dbinom{\chi(a)}{3}=\sum_a \dfrac{\chi(a) (\chi(a)-1) (\chi(a)-2)}{6}$

である (この値を $F(l,r)$ とおく).

$F(l,r)$ と $F(l \pm 1, r)$ や $F(l, r \pm 1)$ の差を調べる.

$F(l-1,r), F(l,r+1)$ のように区間が広がる場合, 新たに加わるのが第 $k$ 項する. このとき, $F(l,r)$ の総和において, 足すべき項が変わるのは $a=A_k$ の場合であるから, $A_l, \dots, A_r$ にある $A_k$ の数を $\beta$ として,

$F(l,r)+\left(\dbinom{\beta+1}{3}-\dbinom{\beta}{3} \right)=F(l,r)+\dfrac{\beta(\beta-1)}{2}$

となる.

$F(l-1,r), F(l,r+1)$ のように区間が狭まる場合も同様に, 区間から外れるのが第 $k$ 項する. このとき, $F(l,r)$ の総和において, 足すべき項が変わるのは $a=A_k$ の場合であるから, $A_l, \dots, A_r$ にある $A_k$ の数を $\beta$ として,

$F(l,r)+\left(\dbinom{\beta-1}{3}-\dbinom{\beta}{3} \right)=F(l,r)-\dfrac{(\beta-1)(\beta-2)}{2}$

となる.

このようにして Mo アルゴリズムによって, うまいこと区間を伸び縮みさせてその差分で結果を更新させていくことにより, $Q$ 個の質問に対して, $O(N \sqrt{Q})$ 時間で求めることができる.