Skip to content
Related Articles

Related Articles

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

Improve Article
Save Article
  • Last Updated : 08 Jul, 2021
Improve Article
Save Article

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)


My Personal Notes arrow_drop_up
Related Articles

Start Your Coding Journey Now!