We set some threshold \(B\) and handle values that appear in \(A\) at least \(B\) times and values that appear less than \(B\) times separately. Here, we call a pair \((i, j)\) such that \(A_i = A_j\) a same-color pair.

• Values that appear at least \(B\) times
  • Since there are at most \(N/B\) such values, even if we perform DFS for each of them in \(O(N)\) time to find the sum of distances of the corresponding same-color pairs, the total time will be \(O(N/B \times N) = O(N^2/B)\).
• Values that appear less than \(B\) times
  • The vertices corresponding to these values form same-color pairs with fewer than \(B\) vertices each, so the total number of same-color pairs for these values is less than \(NB\). Therefore, even if we spend \(O(1)\) time to find the distance for each, the total time will be \(O(NB)\). (Refer to the Wikipedia articles on Lowest common ancestor and Range minimum query. It is not necessarily required to perform the preprocessing in linear time or linear space.)

Setting \(B = \sqrt N\) yields an \(O(N^2/B + NB) = O(N \sqrt N)\) time solution. (This may not be fast enough for some languages.)