Finding pairs (Çözgüdiň düşündirişi)

The condition d(i, j) = d(i, 1) - d(j, 1) holds for every pair of vertices $(i,j)$ where $j$ is an ancestor of $i$ in the depth-first search starting from vertex $1$.

Let us consider an example:

• d(4, 1) = 2, d(4, 1) - d(1, 1) = 2 - 0 = 2;

• d(6, 3) = 2, d(6, 1) - d(3, 1) = 3 - 1 = 2;

• d(2, 2) = 0, d(2, 1) - d(2, 1) = 1 - 1 = 0;

• d(6, 1) = 3, d(6, 1) - d(1, 1) = 3 - 0 = 3;

Let us define the depth $h[v]$ of a vertex $v$ as the number of vertices on the path from the root (vertex $1$) to $v$. In the diagram, the depth is written next to each vertex.

Now, fix a vertex $i$ and ask the question: how many vertices $j$ are there such that the condition d(i, j) = d(i, 1) - d(j, 1) is satisfied?

For example, for i = 6, there are $4$ such vertices: j = 1, 3, 4, 6. Note that h[6] = 4. Thus, for a fixed vertex $i$, there are exactly $h[i]$ suitable vertices $j$.

To solve the problem, we need to compute the depth $h[v]$ for each vertex $v$ in the tree, and then find the sum of all these depths across the tree.

Example

Let us consider the graph from the example. Next to each vertex, we write its depth.

The sum of the depths of all vertices is 1 + 2 + 3 + 3 + 4 = 13.

Algorithm implementation

Define the graph's adjacency list as $g$. The depths of the vertices are stored in the array $h$.

vector<vector<int> > g;
vector<int> h;

The dfs function computes the depth of each vertex $v$. The parent of vertex $v$ is the vertex $p$.

int dfs(int v, int p)
{
  for (int to : g[v])

Consider the edge (v, to). If the vertex $to$ is not the parent of vertex $v$, compute its depth and perform a depth first search starting from it.

    if (to != p)
    {
      h[to] = h[v] + 1;
      dfs(to, v);
    }
  return h[v];
}

The main part of the program. Read the number of vertices $n$ in the tree.

scanf("%d", &n);
g.resize(n + 1);

Construct a tree.

for (i = 2; i <= n; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}

Set the depth of vertex $1$ to be $1$.

h.resize(n + 1);
h[1] = 1;

Start a depth first search from the vertex $1$.

dfs(1, -1);

Compute the answer — the sum of the depths of all vertices.

res = 0;
for (i = 1; i <= n; i++)
  res += h[i];

Print the answer.

printf("%lld\n", res);