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++
// 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 edgesvoid 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 Searchvoid 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 Svoid 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 queriesvoid 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 Codeint 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);} |
Java
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
# Python3 program for the above approachfrom bisect import bisect_left, bisect_righttin = [0] * 100tout = [0] * 100depth = [0] * 100t = 0# Function to add edgesdef Add_edge(parent, child, adj): adj[parent].append(child) adj[child].append(parent) return adj# Function to perform Depth First Searchdef 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 Sdef 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 queriesdef 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 Codeif __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#
// C# equivalent of the above codeusing 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); } }} |
Javascript
let tin = [];let tout = [];let depth = [];let t = 0;// Function to add edgesfunction addEdge(parent, child, adj) { adj[parent].push(child); adj[child].push(parent);}// Function to perform Depth First Searchfunction 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 Sfunction 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 queriesfunction 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 Codelet N = 6;let Q = [[2, 1], [1, 1]];let M = Q.length;numberOfNodesUtil(Q, M, N); |
Output
3 2
Time Complexity: O(N + M*log(N))
Auxiliary Space: O(N)
Feeling lost in the world of random DSA topics, wasting time without progress? It’s time for a change! Join our DSA course, where we’ll guide you on an exciting journey to master DSA efficiently and on schedule.
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
