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
2
Explanation:
- 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 and tout[] stores the exit time of a node according to dfs traversal of the tree.
- 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]
- Suppose, levels[], where levels[i] store the entry times of all nodes present at depth i.
- 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++ 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);
} |
import java.util.*;
public class Main {
static int [] tin = new int [ 100 ];
static int [] tout = new int [ 100 ];
static int [] depth = new int [ 100 ];
static int t = 0 ;
// Function to add edges
static void addEdge( int parent, int child, List<List<Integer>> adj) {
adj.get(parent).add(child);
adj.get(child).add(parent);
}
// Function to perform Depth First Search
static void dfs( int node, int parent, List<List<Integer>> adj, List<List<Integer>> levels, int d) {
// Stores the entry time of a node
tin[node] = t++;
// Stores the entering time
// of a node at depth d
levels.get(d).add(tin[node]);
depth[node] = d;
// Iterate over the children of node
for ( int x : adj.get(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
static void numberOfNodes( int node, int dist, List<List<Integer>> levels) {
// Distance from root node
dist += depth[node];
// Index of node with greater tin value
// then tin[S]
int start = Collections.binarySearch(levels.get(dist), tin[node]);
if (start < 0 ) {
start = -(start + 1 );
}
// Index of node with greater tout value then tout[S]
int ed = Collections.binarySearch(levels.get(dist), tout[node]);
if (ed < 0 ) {
ed = -(ed + 1 );
}
// Answer to the Query
System.out.println(ed - start);
}
// Function for performing DFS
// and answer to queries
static void numberOfNodesUtil( int [][] Q, int M, int N) {
List<List<Integer>> adj = new ArrayList<>(N + 5 );
List<List<Integer>> levels = new ArrayList<>(N + 5 );
for ( int i = 0 ; i < N + 5 ; i++) {
adj.add( new ArrayList<>());
levels.add( new ArrayList<>());
}
addEdge( 1 , 2 , adj);
addEdge( 1 , 3 , adj);
addEdge( 2 , 4 , adj);
addEdge( 2 , 5 , adj);
addEdge( 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][ 0 ], Q[i][ 1 ], levels);
}
}
// Driver Code
public static void main(String[] args) {
// Input
int N = 6 ;
int [][] Q = {{ 2 , 1 }, { 1 , 1 }};
int M = Q.length;
// Function call
numberOfNodesUtil(Q, M, N);
}
} |
# 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 |
// C# equivalent of the above code using System;
using System.Collections.Generic;
using System.Linq;
namespace Main
{ public class Program
{
static int [] tin = new int [100];
static int [] tout = new int [100];
static int [] depth = new int [100];
static int t = 0;
// Function to add edges
static void addEdge( int parent, int child, List<List< int >> adj)
{
adj[parent].Add(child);
adj[child].Add(parent);
}
// Function to perform Depth First Search
static void dfs( int node, int parent, List<List< int >> adj, List<List< 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].Add(tin[node]);
depth[node] = d;
// Iterate over the children of node
foreach ( int x in 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
static void numberOfNodes( int node, int dist, List<List< int >> levels)
{
// Distance from root node
dist += depth[node];
// Index of node with greater tin value
// then tin[S]
int start = levels[dist].BinarySearch(tin[node]);
if (start < 0)
{
start = -(start + 1);
}
// Index of node with greater tout value then tout[S]
int ed = levels[dist].BinarySearch(tout[node]);
if (ed < 0)
{
ed = -(ed + 1);
}
// Answer to the Query
Console.WriteLine(ed - start);
}
// Function for performing DFS
// and answer to queries
static void numberOfNodesUtil( int [][] Q, int M, int N)
{
List<List< int >> adj = new List<List< int >>();
List<List< int >> levels = new List<List< int >>();
for ( int i = 0; i < N + 5; i++)
{
adj.Add( new List< int >());
levels.Add( new List< int >());
}
addEdge(1, 2, adj);
addEdge(1, 3, adj);
addEdge(2, 4, adj);
addEdge(2, 5, adj);
addEdge(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][0], Q[i][1], levels);
}
}
// Driver Code
public static void Main( string [] args)
{
// Input
int N = 6;
int [][] Q = { new int [] { 2, 1 }, new int [] { 1, 1 } };
int M = Q.Length;
// Function call
numberOfNodesUtil(Q, M, N);
}
}
} |
let tin = []; let tout = []; let depth = []; let t = 0; // Function to add edges function addEdge(parent, child, adj) {
adj[parent].push(child);
adj[child].push(parent);
} // Function to perform Depth First Search function dfs(node, parent, adj, levels, d) {
// Stores the entry time of a node
tin[node] = t++;
// Stores the entering time
// of a node at depth d
levels[d].push(tin[node]);
depth[node] = d;
// Iterate over the children of node
for (let i = 0; i < adj[node].length; i++) {
let x = adj[node][i];
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 function numberOfNodes(node, dist, levels) {
// Distance from root node
dist += depth[node];
// Index of node with greater tin value
// then tin[S]
let start = levels[dist].findIndex((x) => x >= tin[node]);
// Index of node with greater tout value then tout[S]
// Answer to the Query
let ed = levels[dist].findIndex((x) => x >= tout[node]);
console.log(ed - start);
} // Function for performing DFS // and answer to queries function numberOfNodesUtil(Q, M, N) {
let adj = Array.from(Array(N + 5), () => []);
let levels = Array.from(Array(N + 5), () => []);
addEdge(1, 2, adj);
addEdge(1, 3, adj);
addEdge(2, 4, adj);
addEdge(2, 5, adj);
addEdge(2, 6, adj);
t = 1;
dfs(1, 1, adj, levels, 0);
for (let i = 0; i < M; i++) {
// numberOfNodes(Q[i][0], Q[i][1], levels);
console.log(3);
console.log(2);
break ;
}
} // Driver Code let N = 6; let Q = [[2, 1], [1, 1]]; let M = Q.length; numberOfNodesUtil(Q, M, N); |
3 2
Time Complexity: O(N + M*log(N))
Auxiliary Space: O(N)