Approach

The approach is the same as the one in the Official Editorial up to the first pseudocode section.

Let’s consider updating it collectively where \(A_i\) are the same. Specifically, during the transition from \(i\), if there exists a \(j\) such that \(A_i=A_j\), then afterwards, sum them all together during the loop of \(j\).

Time Complexity

The above approach can be done in \(O(N\sqrt N)\) time.

An intuitive explanation of the computational complexity is that the worst case scenario is when there are no updates are done “together” and the \(A_i\) values are as small as possible. In this case, \(A\) will look like \(1,\ 2,\ 2,\ 3,\ 3,\ 3,\ 4,\ \cdots\), and here the number of times each is called on the \(j\) side is \(O(\sqrt N)\).

More rigorously, if we define \(G_{k,l} = {\ i \ |\ 0\le i \lt N,\ A_i=k,\ i \bmod k = l }\), and \(f(G_{k,l})=\displaystyle\frac{N}{k}\), then the number of calculations in our loops can be bounded by \(\displaystyle\sum_{G_{k,l}\neq \phi}f(G_{k,l})\). Considering the largest \(N\) values of \(f\), we arrive at the same conclusion as above.

Example Code

AC Code （Python）

In the code, the part to be calculated “together” later is stored in the array Y.