マージテクを用いて解くことができます．

適当な頂点を根とすることで根付き木とします．\(\mathrm{dist}(u,v)=\mathrm{dist}(u,\mathrm{lca}(u,v))+\mathrm{dist}(v,\mathrm{lca}(u,v))\) と分解し，LCA ごとに答えへの寄与を求めることにします．

頂点 \(p\) の相異なる子 \(c_1,c_2\) に対し，\(c_1\) の部分木に含まれる頂点のうち \(A_i=x\) であるような頂点の集合を \(S_{x}\)，\(c_2\) の部分木に含まれる頂点のうち \(A_i=x\) であるような頂点の集合を \(T_{x}\) とします．頂点 \(v\) の根からの深さを \(\mathrm{depth}[v]\) とするとき

\[\begin{aligned} \sum_{u \in S_x} \sum_{v \in T_x} f(u,v)&= \sum_{u \in S_x} \sum_{v \in T_x} \left(\mathrm{dist}(u,p)+\mathrm{dist}(v,p)\right)\\\\ &=\left(\sum_{u \in S_x}\mathrm{dist}(u,p) \right)\times|T_x|+\left(\sum_{v \in T_x} \mathrm{dist}(v,p)\right)\times|S_x|\\\\ &=\left(\sum_{u \in S_x}(\mathrm{depth}[v]-\mathrm{depth}[p])\right)\times|T_x|+\left(\sum_{v \in T_x} (\mathrm{depth}[u]-\mathrm{depth}[p])\right)\times|S_x|\\\\ &=\left(\left(\sum_{u \in S_x}\mathrm{depth}[v]\right)-\mathrm{depth}[p]\times|S_x|\right)\times|T_x|+\left(\left(\sum_{v \in T_x} \mathrm{depth}[u]\right)-\mathrm{depth}[p]\times|T_x|\right)\times|S_x| \end{aligned}\]

となるので \(|S_x|,\displaystyle \sum_{v\in S_x}\mathrm{depth}[v],|T_x|,\sum_{u\in T_x}\mathrm{depth}[u]\) からこの値を求めることができます．

また，直接の子孫関係にあるものに対しても，頂点 \(p\) の部分木に含まれる頂点のうち \(A_i=x\) であるような頂点の集合を \(S_x\) とすると，\(A_p=x\) であった場合の答えへの寄与は

\[\begin{aligned} \sum_{v\in S_x}f(v,p)=\left(\sum_{v\in S_x}\mathrm{depth}[v]\right)-\mathrm{depth}[p]\times |S_x| \end{aligned}\]

として求めることができます．

よって，各頂点に対して，その頂点の部分木に含まれる頂点であって，\(A_i=x\) であるようなものがいくつあるか及びそれらの深さの総和を各 \(x\) についてマージテクしながら持つことでこの問題を解くことができます．

実装例(Python)