Queries to find the count of shortest paths in a Tree that contains a given edge

Given a tree with N vertices numbered from 0 to N – 1, M edges, and Q queries of the form {U, V}, such that there is a direct edge between U and V in the tree. The task for each query is to find all possible shortest paths between any possible unordered pair of vertices from the given tree which contains the edge between the given pair of nodes.

Examples:

Input: N = 6, M[] ={{0, 1}, {0, 2}, {1, 3}, {3, 4}, {3, 5}}, queries[] = {{1, 3}, {0, 2}}
                  0
                /  \
             1     2
           /
         3
       /  \
    4     5
Output:
9
5
Explanation:
Query 1: The edge (1, 3) lies in the paths {1, 3), (1, 3, 4), (1, 3, 5), (0, 3), (0, 4), (0, 5), (2, 3), (2, 4) and (2, 5).
Query 2: The edge (0, 2) lies in the paths (2, 0), (2, 1), (2, 3), (2, 4) and (2, 5).

Input: N = 6, M[] ={{0, 1}, {0, 2}, {2, 3}, {1, 4}, {1, 5}}, queries[] = {{1, 5}, {0, 2}}
                  0
                /  \
             1     2
           /  \    /
        4     5  3
Output:
5
8

Approach: The problem can be solved based on the observation that for any query {U, V}, the shortest path between any pair of nodes in the tree will contain the given edge (U, V) if one of the nodes lie in the subtree of U and the other node lies in the remaining tree. Therefore, the required number of pairs will be:



Count of shortest paths containing (U, V) as an edge = subtreeSize(U) * (N – subtreeSize(V)). 

Therefore, follow the steps below to solve the problem:

  • Perform Depth First Search traversal on the tree starting from the root node.
  • For each node, store the count of nodes in its subtree (includes the node).
  • Iterate over each query (U, V) and calculate:

min( subtreeSize(U), subtreeSize(V)) * ( N – min( subtreeSize(U), subtreeSize(V)) )

Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation for the above approach
 
#include <bits/stdc++.h>
using namespace std;
 
const int sz = 1e5;
 
// Adjacency list to
// represent the tree
vector<int> tree[sz];
 
// Number of vertices
int n;
 
// Mark visited/ unvisited
// vertices
bool vis[sz];
 
// Stores the subtree size
// of the corresponding nodes
int subtreeSize[sz];
 
// Function to create an
// edge between two vertices
void addEdge(int a, int b)
{
    // Add a to b's list
    tree[a].push_back(b);
 
    // Add b to a's list
    tree[b].push_back(a);
}
 
// Function to perform DFS
void dfs(int x)
{
    // Mark the vertex
    // visited
    vis[x] = true;
 
    // Include the node in
    // the subtree
    subtreeSize[x] = 1;
 
    // Traverse all its children
    for (auto i : tree[x]) {
        if (!vis[i]) {
            dfs(i);
            subtreeSize[x]
                += subtreeSize[i];
        }
    }
}
 
// Function to print the
// required number of paths
void countPairs(int a, int b)
{
    int sub = min(subtreeSize[a],
                  subtreeSize[b]);
 
    cout << sub * (n - sub)
         << endl;
}
 
// Driver Code
int main()
{
    // Number of vertices
    n = 6;
 
    addEdge(0, 1);
    addEdge(0, 2);
    addEdge(1, 3);
    addEdge(3, 4);
    addEdge(3, 5);
 
    // Calling modified dfs function
    dfs(0);
 
    // Count pairs of vertices in the tree
    countPairs(1, 3);
    countPairs(0, 2);
    return 0;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation for
// the above approach
import java.util.*;
class GFG{
 
static int sz = (int) 1e5;
 
// Adjacency list to
// represent the tree
static Vector<Integer> []tree = new Vector[sz];
 
// Number of vertices
static int n;
 
// Mark visited/ unvisited
// vertices
static boolean []vis = new boolean[sz];
 
// Stores the subtree size
// of the corresponding nodes
static int []subtreeSize = new int[sz];
 
// Function to create an
// edge between two vertices
static void addEdge(int a, int b)
{
  // Add a to b's list
  tree[a].add(b);
 
  // Add b to a's list
  tree[b].add(a);
}
 
// Function to perform DFS
static void dfs(int x)
{
  // Mark the vertex
  // visited
  vis[x] = true;
 
  // Include the node in
  // the subtree
  subtreeSize[x] = 1;
 
  // Traverse all its children
  for (int i : tree[x])
  {
    if (!vis[i])
    {
      dfs(i);
      subtreeSize[x] += subtreeSize[i];
    }
  }
}
 
// Function to print the
// required number of paths
static void countPairs(int a, int b)
{
  int sub = Math.min(subtreeSize[a],
                     subtreeSize[b]);
 
  System.out.print(sub * (n - sub) + "\n");
}
 
// Driver Code
public static void main(String[] args)
{
  // Number of vertices
  n = 6;
  for (int i = 0; i < tree.length; i++)
    tree[i] = new Vector<Integer>();
  addEdge(0, 1);
  addEdge(0, 2);
  addEdge(1, 3);
  addEdge(3, 4);
  addEdge(3, 5);
 
  // Calling modified dfs function
  dfs(0);
 
  // Count pairs of vertices in the tree
  countPairs(1, 3);
  countPairs(0, 2);
}
}
 
// This code is contributed by 29AjayKumar

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# Python3 implementation for
# the above approach
sz = 100000
 
# Adjacency list to
# represent the tree
tree = [[] for i in range(sz)]
 
# Number of vertices
n = 0
 
# Mark visited/ unvisited
# vertices
vis = [False] * sz
 
# Stores the subtree size
# of the corresponding nodes
subtreeSize = [0 for i in range(sz)]
 
# Function to create an
# edge between two vertices
def addEdge(a, b):
     
    global tree
     
    # Add a to b's list
    tree[a].append(b)
 
    # Add b to a's list
    tree[b].append(a)
 
# Function to perform DFS
def dfs(x):
     
    # Mark the vertex
    # visited
    global vis
    global subtreeSize
    global tree
    vis[x] = True
 
    # Include the node in
    # the subtree
    subtreeSize[x] = 1
 
    # Traverse all its children
    for i in tree[x]:
        if (vis[i] == False):
            dfs(i)
            subtreeSize[x] += subtreeSize[i]
 
# Function to print the
# required number of paths
def countPairs(a, b):
     
    global subtreeSize
    sub = min(subtreeSize[a],
              subtreeSize[b])
 
    print(sub * (n - sub))
 
# Driver Code
if __name__ == '__main__':
     
    # Number of vertices
    n = 6
     
    addEdge(0, 1)
    addEdge(0, 2)
    addEdge(1, 3)
    addEdge(3, 4)
    addEdge(3, 5)
 
    # Calling modified dfs function
    dfs(0)
 
    # Count pairs of vertices in the tree
    countPairs(1, 3)
    countPairs(0, 2)
 
# This code is contributed by SURENDRA_GANGWAR

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// C# implementation for
// the above approach
using System;
using System.Collections.Generic;
class GFG{
   
static int sz = (int) 1e5;
 
// Adjacency list to
// represent the tree
static List<int> []tree =
            new List<int>[sz];
 
// Number of vertices
static int n;
 
// Mark visited/ unvisited
// vertices
static bool []vis = new bool[sz];
 
// Stores the subtree size
// of the corresponding nodes
static int []subtreeSize = new int[sz];
 
// Function to create an
// edge between two vertices
static void addEdge(int a, int b)
{
  // Add a to b's list
  tree[a].Add(b);
 
  // Add b to a's list
  tree[b].Add(a);
}
 
// Function to perform DFS
static void dfs(int x)
{
  // Mark the vertex
  // visited
  vis[x] = true;
 
  // Include the node in
  // the subtree
  subtreeSize[x] = 1;
 
  // Traverse all its children
  foreach (int i in tree[x])
  {
    if (!vis[i])
    {
      dfs(i);
      subtreeSize[x] += subtreeSize[i];
    }
  }
}
 
// Function to print the
// required number of paths
static void countPairs(int a, int b)
{
  int sub = Math.Min(subtreeSize[a],
                     subtreeSize[b]);
 
  Console.Write(sub * (n - sub) + "\n");
}
 
// Driver Code
public static void Main(String[] args)
{
  // Number of vertices
  n = 6;
  for (int i = 0; i < tree.Length; i++)
    tree[i] = new List<int>();
  addEdge(0, 1);
  addEdge(0, 2);
  addEdge(1, 3);
  addEdge(3, 4);
  addEdge(3, 5);
 
  // Calling modified dfs function
  dfs(0);
 
  // Count pairs of vertices in the tree
  countPairs(1, 3);
  countPairs(0, 2);
}
}
 
// This code is contributed by 29AjayKumar

chevron_right


Output: 

9
5









 

Time Complexity: O(N + M + Q)
Auxiliary Space: O(N) 

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.




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.