適当にしきい値 \(B\) を定め、\(A\) に \(B\) 回以上登場する値と \(B\) 回未満しか登場しない値を別に処理します。以下、\(A_i = A_j\) であるような組 \((i, j)\) を同色ペアと呼びます。

• \(B\) 回以上登場する値
  • このような値は \(N/B\) 個以下しか存在しないため、それぞれに対して \(O(N)\) 時間かけて DFS を行って該当する同色ペアの距離の合計を求めても、合計で \(O(N/B \times N) = O(N^2/B)\) 時間で済みます。
• \(B\) 回未満しか登場しない値
  • このような値に対する頂点は、それぞれ \(B\) 個未満の頂点としか同色ペアを形成しないため、このような値に対する同色ペアは合計で \(NB\) 個未満しか存在しません。よって、それぞれに対して \(O(1)\) 時間かけて距離を求めても、合計で \(O(NB)\) 時間で済みます。（Wikipedia の Lowest common ancestor や Range minimum query の記事を参考にしてください。必ずしも前計算を線形時間や線形空間で行う必要はないでしょう。）

\(B = \sqrt N\) とすると \(O(N^2/B + NB) = O(N \sqrt N)\) 時間の解法が得られます。（ある程度言語を選びます。）