Queries for M-th node in the DFS of subtree

Given a tree of N nodes and N-1 edges. Also given an integer M and a node, the task is to print the M-th node in the DFS of the subtree of a given node for multiple queries.

Note: M will not be greater than the number of nodes in the subtree of the given node.

Input: M = 3, node = 1
Output: 4
In the above example if 1 is given as the node, then the DFS of subtree will be 1 2 4 6 7 5 3, hence if M is 3, then the 3rd node is 4

Input: M = 4, node = 2
Output: 7
If 2 is given as the node, then the DFS of the subtree will be 2 4 6 7 5., hence if M is 4 then the 4th node is 7.

Recommended: Please try your approach on {IDE} first, before moving on to the solution.

Approach:

Add the edges between the nodes in an adjacency list.
Call DFS function to generate the DFS of the complete tree.
Use an under[] array to store the height of the subtree under the given node including the node.
In the DFS function, keep incrementing the size of subtree on every recursive call.
Mark the node index in the DFS of complete using hashing.
Let index of given node in the DFS of the tree be ind, then the M-th node will be at index ind + M -1 as the DFS of a subtree of a node will always be a contiguous subarray starting from the node.

Below is the implementation of the above approach.

// C++ program for Queries 
// for DFS of subtree of a node in a tree 
#include <bits/stdc++.h> 
using namespace std; 
const int N = 100000; 
  
// Adjaceny list to store the 
// tree nodes connection 
vector<int> v[N]; 
  
// stores the index of node in DFS 
unordered_map<int, int> mp; 
  
// stores the index of node in 
// original node 
vector<int> a; 
  
// Function to call DFS and count nodes 
// under that subtree 
void dfs(int under[], int child, int parent) 
{ 
  
    // stores the DFS of tree 
    a.push_back(child); 
  
    // hieght of subtree 
    under[child] = 1; 
  
    // iterate for children 
    for (auto it : v[child]) { 
  
        // if not equal to parent 
        // so that it does not traverse back 
        if (it != parent) { 
  
            // call DFS for subtree 
            dfs(under, it, child); 
  
            // add the height 
            under[child] += under[it]; 
        } 
    } 
} 
  
// Function to return the DFS of subtree of node 
int printnodeDFSofSubtree(int node, int under[], int m) 
{ 
    // index of node in the original DFS 
    int ind = mp[node]; 
  
    // height of subtree of node 
    return a[ind + m - 1]; 
} 
  
// Function to add edges to a tree 
void addEdge(int x, int y) 
{ 
    v[x].push_back(y); 
    v[y].push_back(x); 
} 
  
// Marks the index of node in original DFS 
void markIndexDfs() 
{ 
    int size = a.size(); 
  
    // marks the index 
    for (int i = 0; i < size; i++) { 
        mp[a[i]] = i; 
    } 
} 
  
// Driver Code 
int main() 
{ 
    int n = 7; 
  
    // add edges of a tree 
    addEdge(1, 2); 
    addEdge(1, 3); 
    addEdge(2, 4); 
    addEdge(2, 5); 
    addEdge(4, 6); 
    addEdge(4, 7); 
  
    // array to store the height of subtree 
    // of every node in a tree 
    int under[n + 1]; 
  
    // Call the function DFS to generate the DFS 
    dfs(under, 1, 0); 
  
    // Function call to mark the index of node 
    markIndexDfs(); 
  
    int m = 3; 
  
    // Query 1 
    cout << printnodeDFSofSubtree(1, under, m) << endl; 
  
    // Query 2 
    m = 4; 
    cout << printnodeDFSofSubtree(2, under, m); 
  
    return 0; 
} 

Time Complexity: O(1), for processing each query.
Auxiliary Space: O(N)

My Personal Notes

Recommended: Please try your approach on {IDE} first, before moving on to the solution.

Recommended Posts: