Tree Rerooting Technique in C++
By DeepanshuNow let’s learn how to use Tree Rerooting Technique in C++.
Problem description
Consider an unweighted tree with N node and N-1 edges and a positive number D, Find how many distinct pairs of vertices are possible in the tree which is exactly at a Distance of D (Assume (x,y) and (y,x) as same pairs).
Constraint
(Given : 1<=N<=100000 and 1<=D<=1000)
Tree
1–>2,3,4
2–>5,6
3–>7
4–>8,9,10
Let us solve the problem for the above tree.
The valid pairs are (1,5) , (1,6) , (1,7) , (1,8) , (1,9) , (1,10) , (2,3) , (2,4) , (3,4) , (5,6) , (8,9) , (8,10) , (9,10). So, there are 13 valid pairs.
Method 1 ( Naive approach )
A simple method is that we will consider each node as root, and apply DFS every time.
#include<bits/stdc++.h>
using namespace std;
int ans = 0;
void DFS(int node, int parent, int distance, vector<int>v[], int D)
{
if(distance == D)
{
// When distance of the node is D, we increment our answer
++ans;
// If we do not return here then we will call DFS for nodes at distance greater than D
// which we don't require, so we return in this condition
return;
}
for(int neighbour : v[node])
{
if(neighbour == parent) continue;
DFS(neighbour, node, distance+1, v, D);
}
return;
}
int main()
{
// number of nodes and value of D
int n = 10 , D = 2;
// adjecency list
vector<int> v[n+1];
// initialization of tree given in the diagram
v[1].push_back(2), v[2].push_back(1);
v[1].push_back(3), v[3].push_back(1);
v[1].push_back(4), v[4].push_back(1);
v[2].push_back(5), v[5].push_back(2);
v[2].push_back(6), v[6].push_back(2);
v[3].push_back(7), v[7].push_back(3);
v[4].push_back(8), v[8].push_back(4);
v[4].push_back(9), v[9].push_back(4);
v[4].push_back(10), v[10].push_back(4);
for(int i=1; i<=10; ++i)
{
// For every node i from 1 to n , we will consider it as root and call DFS
DFS(i, 0, 0, v, D);
}
// We have calculated each valid pair 2 times
// For example, for the given tree (5,7) is a valid pair
// When we call DFS with 5 as root, we increment answer when 7 is found
// And when we call DFS with 7 as root, we increment answer when 5 is found
// So, we have to divide the obtained answer by 2 to get unique pairs
ans = ans/2;
cout << ans;
return 0;
}
In this method, we are calling the DFS function for each node, and every call will perform O(N) operations in the worst case. Also, the adjacency list will take O(N) space.
So, Time Complexity: O(N2), and Space Complexity is O(N).
Method 2( Using Tree Re-Rooting)
In this method, first, we will assume a node as the root. Let node 1 be the root. Now, we will use DP on Trees to calculate a number of vertices in the subtree of every node at a distance of 1,2,…D, and store it in dp[i][j], where the value of dp[i][j] is equal to a number of vertices in the subtree of i which are at a distance of j. We observe that whichever node we consider as root, we will get the correct answer corresponding to it. For example, here, we initially take 1 as root, and we get dp[1][D]=6, which is equal to a total number of vertices at distance D from 1. So, by this observation, we can conclude that if we make each node as a root and add corresponding values each time, we will get our final answer.
Re-Rooting Method: Our current root is 1, we will now make 2 as root. To do this, we will perform the following steps.
- Subtract contribution of node 2 from node 1.
- Add contribution of node 1 to node 2.
We can observe that if a vertex (common to a subtree of node 1 and node 2), is at distance x from node 1, then the distance from node 2 will be x-1. So to perform step 1, we will subtract all vertices at distance x-1 from node 2, from, vertices at distance x from node 1,(by varying x from 1 to D). Now performing Step 2. We can observe that if a vertex (not in the subtree of node 2), is at distance x from node 2, then its distance from node 1 will be x-1. So, we will add all vertices at distance x-1 from node 1, too, vertices at distance x from node 2,(by varying x from 1 to D). By performing the above steps, we have now made node 2 as root. We will now add the corresponding value obtained (dp[2][D]) to our answer.
In this way, we will make each node a root and add obtained value to get the final answer.
Below is the implementation of the above idea.
#include<bits/stdc++.h>
using namespace std;
int ans=0;
void DFS(int node, int parent, int D, vector<int>v[], vector<int>dp[])
{
// Number of vertices at distance 0 is 1 (i.e. the node itself)
dp[node][0] = 1;
for(int neighbour : v[node]) {
if(neighbour == parent) continue;
DFS(neighbour, node, D, v, dp);
// Number of vertices at distance x in subtree of a
// node is equal to sum of the number of vertices
// at distance x-1 in subtree of its neighbours
for(int x=1; x<=D; ++x) dp[node][x] += dp[neighbour][x-1];
}
}
void ReRoot(int node, int parent, int D, vector<int>v[], vector<int>dp[])
{
// Current node is the root , so we will add dp[node][D] to our answer
ans += dp[node][D];
for(auto neighbour : v[node]){
if(neighbour == parent) continue;
// We will make neighbour as root
// Subtracting contribuion of neighbour from node
// Observe that we are subtracting vertices at distance x-1 from neighbour,
// from the vertices at distance x from node
for(int x=1; x<=D; ++x) dp[node][x] -= dp[neighbour][x-1];
// Adding contribution of node to neighbour
// Observe that we are adding vertices at distance x-1 from node,
// to the vertices at distance x from neighbour
for(int x=1; x<=D; ++x) dp[neighbour][x] += dp[node][x-1];
// By performing above to steps, we have now made neighbour as a root
// So we will recursively call ReRoot function considering neighbour as root
ReRoot(neighbour, node, D, v, dp);
// Now, we have calculated answer for only one neighbour of tne node.
// To calcute answer for other neighbours, we will make node as root again
// To do this, we will perform opposite steps as that from above
// We will first subtract contribution of node from neighbour
for(int x=1; x<=D; ++x) dp[neighbour][x] -= dp[node][x-1];
// Now we will add contribution of neighbour to node
for(int x=1; x<=D; ++x) dp[node][x] += dp[neighbour][x-1];
// After performing above steps, we have made node as root again
}
}
int main()
{
// number of nodes and value of D
int n = 10 , D = 2;
// adjecency list
vector<int> v[n+1];
// initialization of tree given in the diagram
v[1].push_back(2), v[2].push_back(1);
v[1].push_back(3), v[3].push_back(1);
v[1].push_back(4), v[4].push_back(1);
v[2].push_back(5), v[5].push_back(2);
v[2].push_back(6), v[6].push_back(2);
v[3].push_back(7), v[7].push_back(3);
v[4].push_back(8), v[8].push_back(4);
v[4].push_back(9), v[9].push_back(4);
v[4].push_back(10), v[10].push_back(4);
vector<int> dp[n+1];
for(int i=1; i<=n; ++i)
{
// In dp[i][j] we will store number of vertices at distance of j in subtree of i
dp[i].resize(D+1);
}
// first we will call DFS assuming 1 as root, and calculate dp values
DFS(1, 0, D, v, dp);
// We will now call ReRoot function which will make
// each node as root and calculate corresponding answer
ReRoot(1, 0, D, v, dp);
// Here also, we have calculated each valid pair 2 times
// So, we have to divide the obtained answer by 2 to get unique pairs
ans = ans/2;
cout << ans;
return 0;
}
In this method, we are calling ReRoot function for every node, and each call takes O(D) operations to change the root of the tree, and for each node, we maintained a vector of size D
Time complexity: O(N*D), and Space complexity: O(N*D).
Hope you understand the implementation and algorithm used! Thank you
Leave a Reply