考え方

公式解説 の最初の擬似コードの部分までの考え方は同じです。

\(A_i\) が同じところをなるべくまとめて更新することを考えましょう。具体的には\(i\) からの遷移中に \(A_i=A_j\) となる \(j\) が存在すれば、それ以降は \(j\) のループのときにまとめて合計すれば良いです。

計算量について

計算量は \(O(N\sqrt N)\) です。

計算量の直感的な説明は、「まとめて」更新されるのがなく、かつなるべく \(A_i\) たちが小さいときが最悪ケースになりますが、そのとき\(A\) は \(1,\ 2,\ 2,\ 3,\ 3,\ 3,\ 4,\ \cdots\) のようになり、このとき \(j\) 側として呼ばれる回数はそれぞれ \(O(\sqrt N)\) であることから分かります。

より厳密には、 \(G_{k,l} = \{\ i \ |\ 0\le i \lt N,\ A_i=k,\ i \bmod k = l \}\) 、 \(f(G_{k,l})=\displaystyle\frac{N}{k}\) とするとループ数は \(\displaystyle\sum_{G_{k,l}\neq \phi}f(G_{k,l})\) で抑えられますが、 \(f\) が大きい方から \(N\) 個を取ることを考えると上と同じ結論が得られます。

コード例

AC コード （Python）

コードでは、まとめて後で計算する分を配列 Y に格納しています。