Related Articles

# Count of Nodes at distance K from S in its subtree for Q queries

• Last Updated : 08 Jul, 2021

Given a tree consisting of N nodes and rooted at node 1, also given an array Q[] of M pairs, where each array element represents a query of the form (S, K). The task is to print the number of nodes at the distance K in the subtree of the node S for each query (S, K) of the array.

Examples:

Input:  Q[] = {{2, 1}, {1, 1}}, Output:
2
Explanation:

1. Query(2, 1): Print 3, as there are 3 nodes 4, 5, and 6 at the distance of 1 in the subtree of node 2.
2. Query(1, 1): Print 2, as there are 2 nodes 2, and 3 at the distance of 1 in the subtree of node 1.

Input: Edges = {{1, 2}, {2, 3}, {3, 4}}, Q[] = {{1, 2}, {2, 2}}
Output: 1 1

Naive Approach: The simplest approach is for each query to run a Depth First Search(DFS) from node S and find all the nodes that are at a distance K from a given node S

Time Complexity: O(N*Q)
Auxiliary Space: O(1)

Efficient Approach: The above approach can be optimized based on the following observations:

1. Suppose, tin[] stores the entry time of every node and tout[] stores the exit time of a node according to dfs traversal of the tree.
2. Then, for two nodes, A and B, B will be in the subtree of A if and only if:
• tin[B]≥tin[A] and tout[B]≤tout[A]
3. Suppose, levels[], where levels[i] store the entry times of all nodes present at depth i.
4. Then, using binary search nodes at a distance K from a node can be found.

Follow the steps below to solve the problem:

• Initialize three arrays, say tin[], tout[], and depth[] to store the entry time, exit time, and depth of a node respectively.
• Initialize two 2D vectors, say adj and levels, to store the adjacency list and entry times of every node at a specific depth.
• Initialize a variable, say t as 1, to keep track of time.
• Define a recursive DFS function, say dfs(node, parent, d), and perform the following steps:
• Assign t to tin[node] and then increment t by 1.
• Push the tin[node] in the vector levels[d] and then assign d to depth[node].
• Iterate over the children of the node and call the recursive function as dfs(X, node, d+1) for every child X.
• After the above steps, assign t to tout[node] and increment t by 1.
• Call the recursive function dfs(1, 1, 0).
• Traverse the array Q[] using the variable i and do the following:
• Store the value of the current array element as S = Q[i].first, and K = Q[i].second.
• Find the count of all the nodes greater than tin[S] in the vector levels[depth[S]+K] and store it in a variable say L.
• Find the count of all the nodes greater than tout[S] in the vector levels[depth[S]+K] and store it in a variable say R.
• Print the value of R-L as the answer to the current query.

Below is the implementation of the above approach:

## C++

 `// C++ program for the above approach``#include ``using` `namespace` `std;` `int` `tin, tout, depth;``int` `t = 0;` `// Function to add edges``void` `Add_edge(``int` `parent, ``int` `child,``              ``vector >& adj)``{``    ``adj[parent].push_back(child);``    ``adj[child].push_back(parent);``}` `// Function to perform Depth First Search``void` `dfs(``int` `node, ``int` `parent, vector >& adj,``         ``vector >& levels, ``int` `d)``{``    ``// Stores the entry time of a node``    ``tin[node] = t++;` `    ``// Stores the entering time``    ``// of a node at depth d``    ``levels[d].push_back(tin[node]);``    ``depth[node] = d;` `    ``// Iterate over the children of node``    ``for` `(``auto` `x : adj[node]) {``        ``if` `(x != parent)``            ``dfs(x, node, adj, levels, d + 1);``    ``}` `    ``// Stores the Exit time of a node``    ``tout[node] = t++;``}` `// Function to find number of nodes``// at distance K from node S in the``// subtree of S``void` `numberOfNodes(``int` `node, ``int` `dist,``                   ``vector >& levels)``{``    ``// Distance from root node``    ``dist += depth[node];` `    ``// Index of node with greater tin value``    ``// then tin[S]``    ``int` `start = lower_bound(levels[dist].begin(),``                            ``levels[dist].end(), tin[node])``                ``- levels[dist].begin();` `    ``// Index of node with greater tout value then tout[S]``    ``int` `ed = lower_bound(levels[dist].begin(),``                         ``levels[dist].end(), tout[node])``             ``- levels[dist].begin();` `    ``// Answer to the Query``    ``cout << ed - start << endl;``}` `// Function for performing DFS``// and answer to queries``void` `numberOfNodesUtil(pair<``int``, ``int``> Q[], ``int` `M, ``int` `N)``{` `    ``vector > adj(N + 5), levels(N + 5);` `    ``Add_edge(1, 2, adj);``    ``Add_edge(1, 3, adj);``    ``Add_edge(2, 4, adj);``    ``Add_edge(2, 5, adj);``    ``Add_edge(2, 6, adj);` `    ``t = 1;` `    ``// DFS function call``    ``dfs(1, 1, adj, levels, 0);` `    ``// Traverse the array Q[]``    ``for` `(``int` `i = 0; i < M; ++i) {``        ``numberOfNodes(Q[i].first, Q[i].second, levels);``    ``}``}` `// Driver Code``int` `main()``{``    ``// Input``    ``int` `N = 6;``    ``pair<``int``, ``int``> Q[] = { { 2, 1 }, { 1, 1 } };``    ``int` `M = ``sizeof``(Q) / ``sizeof``(Q);` `    ``// Function call``    ``numberOfNodesUtil(Q, M, N);``}`

## Python3

 `# Python3 program for the above approach``from` `bisect ``import` `bisect_left, bisect_right` `tin ``=` `[``0``] ``*` `100``tout ``=` `[``0``] ``*` `100``depth ``=` `[``0``] ``*` `100``t ``=` `0` `# Function to add edges``def` `Add_edge(parent, child, adj):``    ` `    ``adj[parent].append(child)``    ``adj[child].append(parent)``    ``return` `adj` `# Function to perform Depth First Search``def` `dfs(node, parent, d):``    ` `    ``global` `tin, tout, depth, adj, levels, t``    ` `    ``# Stores the entry time of a node``    ``tin[node] ``=` `t``    ``t ``+``=` `1` `    ``# Stores the entering time``    ``# of a node at depth d``    ``levels[d].append(tin[node])``    ``depth[node] ``=` `d` `    ``# Iterate over the children of node``    ``for` `x ``in` `adj[node]:``        ``if` `(x !``=` `parent):``            ``dfs(x, node, d ``+` `1``)` `    ``# Stores the Exit time of a node``    ``tout[node] ``=` `t``    ``t ``+``=` `1` `# Function to find number of nodes``# at distance K from node S in the``# subtree of S``def` `numberOfNodes(node, dist):``    ` `    ``global` `levels, tin, tout``    ` `    ``# Distance from root node``    ``dist ``+``=` `depth[node]` `    ``# Index of node with greater tin value``    ``# then tin[S]``    ``start ``=` `bisect_left(levels[dist], tin[node])` `    ``# Index of node with greater tout value then tout[S]``    ``ed ``=` `bisect_left(levels[dist], tout[node])` `    ``# Answer to the Query``    ``print``(ed ``-` `start)` `# Function for performing DFS``# and answer to queries``def` `numberOfNodesUtil(Q, M, N):``    ` `    ``global` `t, adj` `    ``adj ``=` `Add_edge(``1``, ``2``, adj)``    ``adj ``=` `Add_edge(``1``, ``3``, adj)``    ``adj ``=` `Add_edge(``2``, ``4``, adj)``    ``adj ``=` `Add_edge(``2``, ``5``, adj)``    ``adj ``=` `Add_edge(``2``, ``6``, adj)` `    ``t ``=` `1` `    ``# DFS function call``    ``dfs(``1``, ``1``, ``0``)` `    ``# Traverse the array Q[]``    ``for` `i ``in` `range``(M):``        ``numberOfNodes(Q[i][``0``], Q[i][``1``])` `# Driver Code``if` `__name__ ``=``=` `'__main__'``:``    ` `    ``# Input``    ``N ``=` `6``    ``Q ``=` `[ [ ``2``, ``1` `], [ ``1``, ``1` `] ]` `    ``M ``=` `len``(Q)` `    ``adj ``=` `[[] ``for` `i ``in` `range``(N``+``5``)]``    ``levels ``=` `[[] ``for` `i ``in` `range``(N ``+` `5``)]` `    ``# Function call``    ``numberOfNodesUtil(Q, M, N)` `# This code is contributed by mohit kumar 29`
Output
```3
2```

Time Complexity: O(N + M*log(N))
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.  To complete your preparation from learning a language to DS Algo and many more,  please refer Complete Interview Preparation Course.

In case you wish to attend live classes with experts, please refer DSA Live Classes for Working Professionals and Competitive Programming Live for Students.

My Personal Notes arrow_drop_up