Find the number of distinct pairs of vertices which have a distance of exactly k in a tree

Given an integer k and a tree with n nodes. The task is to count the number of distinct pairs of vertices which have a distance of exactly k.

Examples:

Input: k = 2

Output: 4

Input: k = 3

Output: 2

Approach: This problem can be solved using dynamic programming. For every vertex v of the tree, we calculate values d[v][lev] (0 <= lev <= k). This value indicates the number of vertices having distance lev from v. Note that d[v][0] = 0.
Then we calculate the answer. For any vertex v number of pairs will be a product of the number of vertices at level j – 1 and level k – j.



Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation of the approach
#include <bits/stdc++.h>
using namespace std;
#define N 5005
  
// To store vertices and value of k
int n, k;
  
vector<int> gr[N];
  
// To store number vertices at a level i
int d[N][505];
  
// To store the final answer
int ans = 0;
  
// Function to add an edge between two nodes
void Add_edge(int x, int y)
{
    gr[x].push_back(y);
    gr[y].push_back(x);
}
  
// Function to find the number of distinct
// pairs of the vertices which have a distance
// of exactly k in a tree
void dfs(int v, int par)
{
    // At level zero vertex itself is counted
    d[v][0] = 1;
    for (auto i : gr[v]) {
        if (i != par) {
            dfs(i, v);
  
            // Count the pair of vertices at 
            // distance k
            for (int j = 1; j <= k; j++)
                ans += d[i][j - 1] * d[v][k - j];
  
            // For all levels count vertices
            for (int j = 1; j <= k; j++)
                d[v][j] += d[i][j - 1];
        }
    }
}
  
// Driver code
int main()
{
    n = 5, k = 2;
  
    // Add edges
    Add_edge(1, 2);
    Add_edge(2, 3);
    Add_edge(3, 4);
    Add_edge(2, 5);
  
    // Function call
    dfs(1, 0);
  
    // Required answer
    cout << ans;
  
    return 0;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation of the approach
import java.util.*;
  
class GFG 
{
    static final int N = 5005;
  
    // To store vertices and value of k
    static int n, k;
  
    static Vector<Integer>[] gr = new Vector[N];
  
    // To store number vertices at a level i
    static int[][] d = new int[N][505];
  
    // To store the final answer
    static int ans = 0;
  
    // Function to add an edge between two nodes
    static void Add_edge(int x, int y)
    {
        gr[x].add(y);
        gr[y].add(x);
    }
  
    // Function to find the number of distinct
    // pairs of the vertices which have a distance
    // of exactly k in a tree
    static void dfs(int v, int par) 
    {
        // At level zero vertex itself is counted
        d[v][0] = 1;
        for (Integer i : gr[v])
        {
            if (i != par)
            {
                dfs(i, v);
  
                // Count the pair of vertices at
                // distance k
                for (int j = 1; j <= k; j++)
                    ans += d[i][j - 1] * d[v][k - j];
  
                // For all levels count vertices
                for (int j = 1; j <= k; j++)
                    d[v][j] += d[i][j - 1];
            }
        }
    }
  
    // Driver code
    public static void main(String[] args)
    {
        n = 5;
        k = 2;
        for (int i = 0; i < N; i++)
            gr[i] = new Vector<Integer>();
          
        // Add edges
        Add_edge(1, 2);
        Add_edge(2, 3);
        Add_edge(3, 4);
        Add_edge(2, 5);
  
        // Function call
        dfs(1, 0);
  
        // Required answer
        System.out.print(ans);
  
    }
}
  
// This code is contributed by PrinciRaj1992

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# Python3 implementation of the approach
N = 5005
  
# To store vertices and value of k
n, k = 0, 0
  
gr = [[] for i in range(N)]
  
# To store number vertices at a level i
d = [[0 for i in range(505)] 
        for i in range(N)]
  
# To store the final answer
ans = 0
  
# Function to add an edge between two nodes
def Add_edge(x, y):
    gr[x].append(y)
    gr[y].append(x)
  
# Function to find the number of distinct
# pairs of the vertices which have a distance
# of exactly k in a tree
def dfs(v, par):
    global ans
      
    # At level zero vertex itself is counted
    d[v][0] = 1
    for i in gr[v]:
        if (i != par):
            dfs(i, v)
  
            # Count the pair of vertices at
            # distance k
            for j in range(1, k + 1):
                ans += d[i][j - 1] * d[v][k - j]
  
            # For all levels count vertices
            for j in range(1, k + 1):
                d[v][j] += d[i][j - 1]
  
# Driver code
n = 5
k = 2
  
# Add edges
Add_edge(1, 2)
Add_edge(2, 3)
Add_edge(3, 4)
Add_edge(2, 5)
  
# Function call
dfs(1, 0)
  
# Required answer
print(ans)
  
# This code is contributed by Mohit Kumar

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// C# implementation of the approach
using System;
using System.Collections.Generic;
  
class GFG 
{
    static readonly int N = 5005;
  
    // To store vertices and value of k
    static int n, k;
  
    static List<int>[] gr = new List<int>[N];
  
    // To store number vertices at a level i
    static int[,] d = new int[N, 505];
  
    // To store the readonly answer
    static int ans = 0;
  
    // Function to add an edge between two nodes
    static void Add_edge(int x, int y)
    {
        gr[x].Add(y);
        gr[y].Add(x);
    }
  
    // Function to find the number of distinct
    // pairs of the vertices which have a distance
    // of exactly k in a tree
    static void dfs(int v, int par) 
    {
        // At level zero vertex itself is counted
        d[v, 0] = 1;
        foreach (int i in gr[v])
        {
            if (i != par)
            {
                dfs(i, v);
  
                // Count the pair of vertices at
                // distance k
                for (int j = 1; j <= k; j++)
                    ans += d[i, j - 1] * d[v, k - j];
  
                // For all levels count vertices
                for (int j = 1; j <= k; j++)
                    d[v, j] += d[i, j - 1];
            }
        }
    }
  
    // Driver code
    public static void Main(String[] args)
    {
        n = 5;
        k = 2;
        for (int i = 0; i < N; i++)
            gr[i] = new List<int>();
          
        // Add edges
        Add_edge(1, 2);
        Add_edge(2, 3);
        Add_edge(3, 4);
        Add_edge(2, 5);
  
        // Function call
        dfs(1, 0);
  
        // Required answer
        Console.Write(ans);
  
    }
}
  
// This code is contributed by Rajput-Ji

chevron_right


PHP

filter_none

edit
close

play_arrow

link
brightness_4
code

<?php
// PHP implementation of the approach
$N = 5005;
  
// To store vertices and value of k
$gr = array_fill(0, $N, array());
  
// To store number vertices 
// at a level i
$d = array_fill(0, $N
     array_fill(0, 505, 0));
  
// To store the final answer
$ans = 0;
  
// Function to add an edge between 
// two nodes
function Add_edge($x, $y)
{
    global $gr;
    array_push($gr[$x], $y);
    array_push($gr[$y], $x);
}
  
// Function to find the number of distinct
// pairs of the vertices which have a 
// distance of exactly k in a tree
function dfs($v, $par)
{
    global $d, $ans, $k, $gr;
      
    // At level zero vertex itself
    // is counted
    $d[$v][0] = 1;
    foreach ($gr[$v] as &$i
    {
        if ($i != $par
        {
            dfs($i, $v);
  
            // Count the pair of vertices 
            // at distance k
            for ($j = 1; $j <= $k; $j++)
                $ans += $d[$i][$j - 1] * 
                        $d[$v][$k - $j];
  
            // For all levels count vertices
            for ($j = 1; $j <= $k; $j++)
                $d[$v][$j] += $d[$i][$j - 1];
        }
    }
}
  
// Driver code
$n = 5;
$k = 2;
  
// Add edges
Add_edge(1, 2);
Add_edge(2, 3);
Add_edge(3, 4);
Add_edge(2, 5);
  
// Function call
dfs(1, 0);
  
// Required answer
echo $ans;
  
// This code is contributed by mits 
?>

chevron_right


Output:

4

competitive-programming-img




My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.