The problem asks us to count pairs of vertices $$$(u,v)$$$ satisfying:

• $$$u \lt v$$$
• $$$C_u=C_v$$$
• $$$\operatorname{dist}(u,v)=D$$$, where $$$D$$$ is the diameter of the tree.

The diameter is the maximum distance between any two vertices in a tree. We can calculate it using two BFS traversals and reconstruct the diameter path using parent pointers.

First, we perform BFS from an arbitrary vertex (vertex $$$1$$$) to find a farthest vertex $$$A$$$. Then, we perform another BFS starting from $$$A$$$ to find a farthest vertex $$$B$$$.

The distance between $$$A$$$ and $$$B$$$ gives the diameter:

We also store the parent array during the second BFS to reconstruct the diameter path.

The center of a tree is the middle of its diameter. Depending on whether the diameter is even or odd, there are two possibilities.

Case 1: Diameter is even

Let the diameter be $$$D=2k$$$. Since the diameter is even, the tree has a unique center vertex $$$c$$$.

We can obtain this center from the reconstructed diameter path:

Consider two vertices $$$u$$$ and $$$v$$$ whose distance equals the diameter. Since their path is a diameter path, it must pass through the center $$$c$$$. Therefore,

Since $$$c$$$ is the center, the maximum distance of any vertex from $$$c$$$ is $$$k$$$. Hence,

For their sum to equal $$$2k$$$, both distances must be exactly $$$k$$$.

Additionally, the two vertices must belong to different connected components after removing $$$c$$$, since their path must pass through the center.

Thus, the problem reduces to counting pairs of vertices that:

• Are at distance exactly $$$k$$$ from the center.
• Belong to different subtrees of the center.
• Have the same color.

We first perform BFS from the center to calculate the distance of every vertex from it.

Next, we process each subtree adjacent to the center independently. We collect all vertices at distance exactly $$$k$$$ from the center in the current subtree.

We maintain a frequency map mp, where:

For every collected vertex $$$u$$$, we query:

After processing all queries for the current subtree, we insert its vertices into the map by incrementing their corresponding color frequencies.

Processing queries before insertions ensures that we only count pairs belonging to different subtrees.

Case 2: Diameter is odd

Let the diameter be $$$D=2k+1$$$. Since the diameter is odd, the tree has two centers, say $$$c_1$$$ and $$$c_2$$$, connected by an edge.

We can obtain these centers from the reconstructed diameter path:

Consider two vertices $$$u$$$ and $$$v$$$ whose distance equals the diameter. Their path must pass through the central edge $$$(c_1,c_2)$$$.

Without loss of generality, let $$$u$$$ belong to the component containing $$$c_1$$$, and $$$v$$$ belong to the component containing $$$c_2$$$, after removing the central edge.

We have:

Since the diameter is $$$2k+1$$$,

The maximum distance of any vertex from either center, within its respective component, is $$$k$$$. Therefore,

Thus, the problem reduces to counting pairs of vertices having the same color, where one vertex is at distance $$$k$$$ from $$$c_1$$$ on its side and the other is at distance $$$k$$$ from $$$c_2$$$ on its side.

We solve this using a frequency map.

Step 1: Process the first component.

We perform BFS starting from $$$c_1$$$, keeping track of the current vertex, its parent, and its distance from $$$c_1$$$.

We explicitly avoid traversing the central edge $$$(c_1,c_2)$$$ so that we remain within the component containing $$$c_1$$$.

Whenever we reach a vertex at distance exactly $$$k$$$, we increment its color frequency in the map:

Step 2: Process the second component.

Similarly, we perform BFS starting from $$$c_2$$$, avoiding the central edge $$$(c_2,c_1)$$$.

Whenever we reach a vertex $$$v$$$ at distance exactly $$$k$$$, we add its corresponding color frequency to the answer:

Since the two components are disjoint, every counted pair automatically has distance exactly $$$2k+1$$$, and every valid pair is counted exactly once.