# Find equidistant triplets in a tree

Given a tree $T$ with $n$ vertices, we want to find the number of triplets of vertices $(a,b,c)$ such $d(a,b) = d(b,c) = d(c,a)$ where $d$ is the distance function (length of the shortest path between two nodes).

It's pretty easy to do it in $O(n^3)$ time. Is it possible to do it faster?
I think that on-line algorithm and pre-processing should help.

• Don't answer this question, this is from codechef February long challenge. codechef.com/FEB20A/problems/CHGORAM2 – Rakesh Pandit Feb 15 at 9:41
• Here is an idea. Perform BFS. Now one can think of partitioning of the vertex set according to the level of the vertex from root(starting point of BFS). Let $V_1,V_2,\ldots, V_k$ is the partitioning of the vertex set $V$. Important point $d(x,y) = d(x_1,y_1)$ where $x_1$ and $x_2$ belong same partition and $y_1$ and $y_2$ belong to the same partition. – Shiv Feb 15 at 13:36
• Also note that vertices $a$ and $c$ must be from the same bag. So after partitioning you need to count $(_,_,_)$ such that first and third are from same partition. Seems doable in linear time. – Shiv Feb 15 at 17:30

Consider such a triplet $(a,b,c)$, and root $T$ at $a$. Let $e$ be the lowest common ancestor of $b$ and $c$. If $e = a$ then $d(b,c) = 2d(a,c)$, and we conclude that $a=b=c$. If $e \neq a$, then if we now root $T$ at $e$ then $e$ is the lowest common ancestor of $a,b,c$ (pairwise). That means that $a,b,c$ belong to different branches below $e$.
This gives rise to the following algorithm. First, there are $n$ triplets of the form $(a,a,a)$, which you may or may not want to count. Second, each triplet $(a,b,c)$ comes with its unique vertex $e$. So we enumerate over the vertices $e$, root $T$ at $e$, and for each child $f$ of $e$, count the number of vertices at any given depth; this takes time $O(n)$. Now for each depth $D$, we are given a list $a_1,\ldots,a_k$ corresponding to the number of vertices at depth $D$ for the children of $e$. We want to compute $$\sum_{i<j<k} a_i a_j a_k = \frac{1}{6} \left(\sum_i a_i\right)^3 - \frac{1}{2} \left(\sum_i a_i^2\right) \left(\sum_i a_i\right) + \frac{1}{3} \sum_i a_i^3.$$ This phase can also be accomplished in $O(n)$ (going over all possible depths), with some care (you need to sort the children according to depth, but that can be done in linear time since the depths are integers from $0$ to $n-1$; other details left to you). Therefore the entire algorithm is $O(n^2)$.
My guess is that you can find the number of triplets in $\tilde{O}(n)$.