Distance in Tree#4104

Distance in Tree (Editorial)

Let $d p [v] [x]$ dp[v][x] denote the number of vertices at distance $x$ from vertex $v$ in the subtree rooted at $v$ . Set $d p [v] [0] = 1$ dp[v][0] = 1, since the only vertex at distance $0$ from $v$ is $v$ itself.

Suppose that vertex $v$ has children $u_{1}, u_{2}, ..., u_{p}$ u_1, u_2, ..., u_p. Then the number of vertices at distance $x$ from vertex $v$ in the subtree rooted at $v$ is equal to the sum of the numbers of vertices at distance $x - 1$ x - 1 from each of the vertices $u_{1}, u_{2}, ..., u_{p}$ u_1, u_2, ..., u_p.

d p [v] [x] = i = 1 \sum p d p [u_{i}] [x - 1]

dp[v][x] = \sum_{i = 1}^p dp[u_i][x-1]

Consider the subtree rooted at vertex $v$ . Let's count the number of paths of length $k$ that lie entirely within this subtree and pass through vertex $v$ . We divide such paths into two groups:

paths that start at vertex $v$ . Their number is equal to $d p [v] [k]$ dp[v][k].
paths that pass through vertex $v$ and whose endpoints are located in the subtrees of the children of vertex $v$ .

Consider a path of length $k$ that passes through the root of the subtree $v$ . Let the distance from vertex $v$ to one end of the path be $x (1 \leq x \leq k - 1)$ x~(1 \le x \le k - 1). Suppose this end lies in the subtree of some child $u_{1}$ u_1. Then the distance from $v$ to the other end of the path is $k - x$ k - x, and this other end can be located in any other subtree except the subtree rooted at $u_{1}$ u_1.

If the distance from the first end of the path to vertex $v$ is $x$ , then the distance from this end to vertex $u_{1}$ u_1 is $x - 1$ x - 1. The number of ways to choose such a first end is $d p [u_{1}] [x - 1]$ dp[u_1][x - 1].

The number of ways to choose the second end of the path is equal to the number of vertices at distance $k - x$ k - x from vertex $v$ that do not belong to the subtree rooted at $u_{1}$ u_1. This number is

d p [v] [k - x] - d p [u_{1}] [k - x - 1]

dp[v][k - x] - dp[u_1][k - x - 1]

According to the multiplication principle, the total number of such paths is

d p [u_{1}] [x - 1] \cdot (d p [v] [k - x] - d p [u_{1}] [k - x - 1])

dp[u_1][x - 1] \cdot (dp[v][k - x] - dp[u_1][k - x - 1])

Thus, the total number of such paths is

d p [v] [k] + 1/2 i = 1 \sum p x = 1 \sum k - 1 d p [u_{i}] [x - 1] \cdot (d p [v] [k - x] - d p [u_{i}] [k - x - 1])

dp[v][k] + 1/2 \sum_{i = 1}^p \sum_{x = 1}^{k - 1} dp[u_i][x-1] \cdot (dp[v][k-x]-dp[u_i][k-x-1])

It remains to sum the above expression over all vertices of the tree.

Example

Consider the tree presented in the example.

Algorithm implementation

The array $g$ stores the tree representation in the form of adjacency lists. The array $d p$ is used to implement dynamic programming.

vector<vector<int> > g;
vector<vector<int> > dp;

The function dfs performs a depth-first search starting from vertex $v$ . The value $p$ denotes the parent of vertex $v$ . During the traversal, the $d p$ array is constructed.

void dfs(int v, int p = -1)
{

Initialize the array $d p [v]$ .

  dp[v].resize(k+1);
  dp[v][0] = 1;

Iterate over all vertices $t o$ that can be reached from vertex $v$ .

  for (int to : g[v])
  {
    if (to == p) continue;

    dfs(to,v);

    for(int x = 1; x <= k; x++)
      dp[v][x] += dp[to][x-1];
  }
}

C++

The function dfsRes performs a depth-first search starting from vertex $v$ and is intended to compute the desired answer.

void dfsRes(int v, int p = -1)
{

  for (int to : g[v])
  {
    if (to == p) continue;

    dfsRes(to,v);

    for(int x = 1; x < k; x++)
      resA += 1LL * dp[to][x-1] * (dp[v][k-x] - dp[to][k-x-1]);
  }

  resB += dp[v][k];
}

C++

The main part of the program. Read the input data. Build the tree.

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

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

C++

Using the first depth-first search, construct the $d p$ array.

dp.resize(n+1);
dfs(1);

Compute and print the required answer.

resA = resB = 0;
dfsRes(1);
res = resA / 2 + resB;
printf("%lld\n",res);

Python implementation

The function dfs performs a depth-first search starting from vertex $v$ . The value $p$ denotes the parent of vertex $v$ . During the traversal, the $d p$ array is constructed.

def dfs(v, p):

Initialize the array $d p [v]$ .

  dp[v] = [0] * (k + 1)
  dp[v][0] = 1

Iterate over all vertices $t o$ that can be reached from vertex $v$ .

  for to in g[v]:
    if to == p: continue
    dfs(to, v)

    for x in range(1, k + 1):
      dp[v][x] += dp[to][x - 1]

The function dfsRes performs a depth-first search starting from vertex $v$ and is intended to compute the desired answer.

def dfsRes(v, p):
  global resA, resB
  for to in g[v]:
    if to == p: continue
    dfsRes(to, v)

    for x in range(1, k):
      resA += dp[to][x - 1] * (dp[v][k - x] - dp[to][k - x - 1])

  resB += dp[v][k]

Python

The main part of the program. Read the number of vertices $n$ of the tree, and the distance between vertices $k$ .

n, k = map(int, input().split())

The list $g$ stores the tree representation using adjacency lists. The list $d p$ is used to implement dynamic programming.

g = [[] for _ in range(n + 1)]
dp = [[] for _ in range(n + 1)]

Read the list of edges. Build the tree.

for _ in range(n - 1):
  a, b = map(int, input().split())
  g[a].append(b)
  g[b].append(a)

Using the first depth-first search, construct the $d p$ array.

dfs(1, -1)

Compute and print the required answer.

resA, resB = 0, 0
dfsRes(1, -1)
res = resA // 2 + resB
print(res)