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

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:3

2Explanation:

- 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.
- 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:

- Suppose,
tin[]stores the entry time of every node andtout[]stores the exit time of a node according to dfs traversal of the tree.- Then, for two nodes,
AandB,Bwill be in the subtree ofAif and only if:

tin[B]â‰¥tin[A] and tout[B]â‰¤tout[A]- Suppose,
levels[], wherelevels[i]store the entry times of all nodes present at depthi.- Then, using binary search nodes at a distance
Kfrom 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.**

- Assign
- 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.

- Store the value of the current array element as

Below is the implementation of the above approach:

## C++

`// C++ program for the above approach` `#include <bits/stdc++.h>` `using` `namespace` `std;` `int` `tin[100], tout[100], depth[100];` `int` `t = 0;` `// Function to add edges` `void` `Add_edge(` `int` `parent, ` `int` `child,` ` ` `vector<vector<` `int` `> >& adj)` `{` ` ` `adj[parent].push_back(child);` ` ` `adj[child].push_back(parent);` `}` `// Function to perform Depth First Search` `void` `dfs(` `int` `node, ` `int` `parent, vector<vector<` `int` `> >& adj,` ` ` `vector<vector<` `int` `> >& 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<vector<` `int` `> >& 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<vector<` `int` `> > 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[0]);` ` ` `// 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**.