# 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++ implementation for the above approach`   `#include ` `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;` `}`

 `// Java implementation for ` `// the above approach` `import` `java.util.*;` `class` `GFG{`   `static` `int` `sz = (``int``) 1e5;`   `// Adjacency list to` `// represent the tree` `static` `Vector []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();` `  ``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`

 `# 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`

 `// 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`

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.

Article Tags :