Last Updated: 18 Sep, 2021
Restore Peace
Hard
Asked in companies
There is a kingdom with ‘N’ cities that are connected by roads. The cities are connected to form a tree, where cities represent the nodes, and the roads represent the edges. The cities are numbered from 0 to ‘N-1’ and the length of each road is 1 unit.
A tree is a hierarchical data structure defined as a collection of nodes.
Now, a battle is going on in some cities of the kingdom, where the king must visit to restore peace. Initially, the king is at the 0th node(root node). The cities where battles are going on will be given as an input array, and the king has to visit in the same order as given in the input array 'battles'. You have to find the total length of the road the king has to travel to restore peace in all the cities.
For Example:
For the test case:
N = 3, Q = 2
Edge 1 connects 0 and 1
Edge 2 connects 0 and 2
battles = [1, 2]
The king has to visit first city 1 and then city 2. To visit city 1, he has to travel 1 unit, and now he is in city 1. To visit city 2 from city 1, he has to travel another 2 units. So overall, he has to travel 3 units of distance.
Input Format :
The first line will contain a single integer ‘T’ denoting the number of test cases. Then ‘T’ test cases follow.
The first line of each test case will contain two space-separated integers, ‘N’ and ‘Q’, denoting the total number of cities and the number of cities the king has to visit, respectively.
The next ‘N - 1’ line of each test case will contain two space-separated integers U[i] and V[i] denoting the two cities each road connects.
The next line of each test case will contain ‘Q’ space-separated integers denoting the cities where the king has to visit.
Output Format :
For each test case, print a single integer denoting the total length of the road the king has to travel to restore peace in all the cities.
Note :
You are not required to print anything; it has already been taken care of. Just implement the function.
Constraints:
1 <= T <= 10
1 <= Q <= N <= 10^5
0 <= U[i], V[i] <= N - 1
It is guaranteed that the sum of ‘N’ and sum of ‘Q’ over all test cases doesn’t exceed 10^5.
Time Limit: 1 sec.
Approaches
We will iterate over the battle cities in the given order and calculate the distance between the current position of the king and his next position each time separately.
To calculate the distance between two cities, we will use the breadth-first search algorithm.
Algorithm:
- bfs function:
- This function will take the current city, “cur” and the next city, “tar” of the king, and the graph as arguments and will return the distance between these two cities.
- Declare two arrays - “vis” and “dist” to store whether a city is visited or not and to store the distance of each city from the “cur” city, respectively.
- Declare a queue of integers, and push the “cur” city to it.
- Set the distance of the “cur” city as 0.
- While ‘q’ is not empty.
- Take out the front element of the ‘q’.
- If this is the “tar” city, return the distance of this city from the “dist” array.
- Iterate over the adjacent cities.
- If the city is not visited
- Mark it as visited
- Push it to the queue
- Update the “dist” array.
- If the city is not visited
- given function:
- Create the graph from the given input.
- Declare two variables - “ans” to store the answer and “cur” to store the current city of the king.
- Iterate over the battle cities.
- Declare a visited array “vis”.
- Calculate the distance between the current city and the next city using the “bfs” function. Add this to the “ans”.
- Update the “cur” as the next city.
- Return the “ans”.
We will use Lowest Common Ancestor(LCA) to find the distance between two cities in logarithmic time. After knowing the LCA of the two cities we can easily calculate the distance between them. Now to calculate the LCA we have to precalculate two arrays - “level” and “LCA” array. “Level” array will store the “level” of each node from the 0th node(root node), and the “LCA” array will store the ancestors of each node in the power of 2. For example, for each node “LCA” will store the 2^0th, 2^1th, 2^2th ancestors, and so on till it exists.
The distance between two nodes A and B will be “level[A]” + “level[B]” - 2 * “level[LCA of A and B]”.
For more information about LCA, follow the given link below:
https://cp-algorithms.com/graph/lca_binary_lifting.html
Algorithm:
- dfs function:
- This function will take the graph, level array, LCA array, node, current level, and parent as arguments and store the level of each node in the “level” array, and the 2^0th ancestor in the “LCA” array.
- This function will take the graph, level array, LCA array, node, current level, and parent as arguments and store the level of each node in the “level” array, and the 2^0th ancestor in the “LCA” array.
- init function:
- This function will take the graph, “LCA” array, number of cities, and “maxN” as arguments and will fill the “LCA” array.
- getLCA function:
- This function will take two cities, the level array, “LCA” array, and “maxN” as arguments and return the LCA of the two cities.
- getDist function:
- This function will take two cities, the level array, “LCA” array, and “maxN” as arguments, and return the distance between the two cities.
- This function will take two cities, the level array, “LCA” array, and “maxN” as arguments, and return the distance between the two cities.
- given function:
- Create the graph from the given input.
- Declare an integer variable “maxN” and initialize it with 20(close to log2(10^5)).
- Declare a vector “level” of size ‘N’.
- Declare a 2d vector of size ‘N’ x “maxN”.
- Call the “dfs” function for precomputation to storing the ancestors of the node.
- Cal the “init” function.
- Declare two variables - “ans” to store the answer and “cur” to store the current city of the king.
- Iterate over the battle cities.
- Calculate the distance between the current city and the next city using the “getDist” function.
- Add this to the “ans”.
- Update the “cur” as the next city.
- Return the “ans”.