PROBLEM LINK:
Author: Sergey Kulik
Tester: Jingbo Shang
Editorialist: Ajay K. Verma
DIFFICULTY:
Medium-Hard
PREREQUISITES:
PROBLEM:
We are given a tree with edge weights belonging to {1, 2}. We need to handle the queries of the following form: given two nodes u and v in the tree, and and integer c, find the minimum number steps needed to reach from u to v, where in each step a maximum distance of c unit can be covered.
EXPLANATION:
We divide the queries into two classes based on the value c. The queries for which c is smaller than N0.5 are kept in the first class, while the remaining queries are kept in the second class.
There is some common pre-processing that we need to do for the queries of both classes. Let us first discuss this pre-processing step.
First, we make the tree rooted, and maintain the following information for each node:
- distance of this node from root (stored in array dist[1…N])
- number of nodes in the path from root to this node (stored in array ht[1…N])
- 2k -th successor of this node, for all 0 <= k <= lg N (stored in array pt[1…N][0…lg N])
The three arrays can be computed easily using a depth first traversal of the tree.
void dfs(int u, int **adj, int **pt, int *dist, int *ht) { // special case for root if (pt[u][0] == -1) { dist[u] = 0; ht[u] = 0; } for (int v : adj[u]) { if (v == pt[u][0]) continue; // update child's data dist[v] = dist[u] + wt(u, v); ht[v] = 1 + ht[u]; pt[v][0] = u; for (int i = 1 to lg N) pt[v][i] = pt[v][i - 1] == -1 ? -1 : pt[pt[v][i - 1]][i - 1]; // recursion dfs(v, adj, pt, dist, ht); } }
Using the above information, the lowest common ancestor of two nodes u and v can be computed easily as shown below:
// Assumes that ht[u] >= ht[v] int lca(int u, int v, int **pt, int *ht) { // Move from u to its ancestor, // until it is at the same level as v for (i = lg N to 0) { if (pt[u][i] == -1) continue; if (ht[pt[u][i]] >= ht[v]) u = pt[u][i]; } // At this point ht[u] and ht[v] must be equal if (u == v) return u; // move simultaneously u and v towards root, // until they merge for (i = lg N to 0) { if (pt[u][i] != pt[v][i]) { u = pt[u][i]; v = pt[v][i]; } } return pt[u][0]; }
In order to handle the queries of first type, we need to some additional pre-processing, while the queries of the second class can be handled without any additional pre-processing.
Let us first discuss how to handle the queries of the second class.
Queries with Large Value of c:
For the second kind of query, compute the lowest common ancestor w of nodes u and v, and then compute the number of steps to reach from u to w, and the number of steps to reach from v to w. This can be done using the pre-processed data computed in the previous section.
// returns the ancestor node of u that will be reached // in a single step from u. int single_jump(int u, int c, int *dist, int **pt) { int st = dist[u]; for (int i = lg N to 0) { int v = pt[u][i]; if (st - dist[v] > c) continue; // move to v, if distance between u and v // does not exceed c u = v; } return u; } // moves up from u to w, returns the following two values: // 1) the number of steps taken to reach the closest // successor of w from u, such that the next step will take // us beyond w. // 2) the distance between w and this closest successor // of w. struct info { int num_steps; int dist_remaining; }; info move_up (int u, int w, int c, int **pt, int *dist, int *ht) { int num_steps = 0; int dist_remaining = 0; while (u != w) { // Use a single step int tmp = single_jump(u, c, dist, pt); if (ht[tmp] > ht[w]) { // we are still below w ++num_steps; u = tmp; } else { // we cannot take this step, otherwise // we will move beyond w dist_remaining = dist[u] - dist[w]; break; } } return info{num_steps, dist_remaining}; }
Note that the complexity of single_jump is O (lg N), while the complexity of move_up depends on the number of steps taken to reach from u to w. Since the value of c is >= N0.5, number of steps cannot be higher that N0.5. Hence, the complexity of move_up is O (N0.5 lg N).
In a similar way, we can compute the number of steps taken to reach from v to w, and the remaining distance that we need to cover to reach w.
I1 = move_up(u, w, c, pt, dist, ht);
I2 = move_up(v, w, c, pt, dist, ht);
Now, we have taken (I1.num_steps + I2.num_steps) steps and we still need to cover (I1.dist_remaining + I2.dist_remaining) distance. This will require ceil ((I1.dist_remaining + I2.dist_remaining) / c) additional steps.
Hence, the queries of the second class can be handled in O (N0.5 lg N) time.
Queries with Small Value of c:
Unfortunately, we cannot use the above approach for handling the queries of first class, as the number of steps to reach from u to w can be as large as O (N).
However, in this class there are only N0.5 distinct values of c, so we can do some preprocessing that can be used to process multiple jumps in O (1) time.
More formally for each value of c, we create an array P[1…N][0…lg N] such that P[i][j] denote the ancestor node of i that can be reached after making 2j jumps from i.
The array P[][] can be computed using a depth first traversal, in the same way as the array pt[][] was computed. Now using the array P, we can compute the number of steps to reach from u to w in O (lg N) time.
info move_up (int u, int w, int c, int **pt, int *dist, int *ht, int **P) { int num_steps = 0; for (int i = lg N to 0) { int tmp = P[u][i]; if (ht[tmp] > ht[w]) { // we are still below w num_steps += 2^i u = tmp; } } return info{num_steps, dist[u] - dist[w]};
Time Complexity:
Preprocessing time: O (N1.5 lg N)
Query time: O (N0.5 lg N)
AUTHOR’S AND TESTER’S SOLUTIONS:
Author’s solution will be put up soon.
Tester’s solution will be put up soon.