Skip to content
Related Articles

Related Articles

Improve Article

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




My Personal Notes arrow_drop_up
Recommended Articles
Page :