PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: raysh07
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Binary search, DFS, dynamic programming
PROBLEM:
There are N cities connected by pipes in the form of a tree. Pipe i connects cities i and p_i, and has a leakiness of A_i.
A pipe with a leakiness of A_i will lose every A_i-th unit of water that passes through it.
For each city i, K units of water will be sent from city 1 to i, via the unique sequence of pipes connecting 1 and i.
This will result in some of the water getting lost before it reaches i; let R_i be the amount of water that reaches i.
You can fix at most one pipe, which will ensure no water leaks through it.
Find the maximum possible value of \min(R_1, \ldots, R_N) you can achieve.
EXPLANATION:
First, note that if X units of water pass through a pipe with leakiness Y, the amount of water that will be lost is exactly \left\lfloor \frac{X}{Y} \right\rfloor.
This is because every Y-th unit of water will be lost, so we end losing exactly one unit for each positive multiple of Y that’s \le X.
Now, let’s compute how much water each city will receive, before any fixes i.e. the initial R_i values.
Because of the sequential nature of the process, this is in fact quite simple:
- R_1 = K, since it doesn’t need to pass through any pipes.
- For each i \gt 1, R_i = R_{p_i} - \left\lfloor \frac{R_{p_i}}{A_i} \right\rfloor
This is because there will be R_{p_i} units of water reaching p_i by definition; and then we just subtract out how much of that will be lost when moving through the pipe connecting i and p_i.
Next, let’s attempt to actually solve the problem.
We’re allowed to fix at most one pipe.
Suppose we fix the pipe connecting i and p_i.
The effect of this would be to set R_i = R_{p_i} (since there would be no loss of water); but then the effect would cascade down to the subtree of i since there’s now more water at each of them - and recalculating everything would take \mathcal{O}(N) time, which is clearly too slow.
Instead, we use a rather standard trick when dealing with objective functions of the form “maximize the minimum of several quantities” - binary search!
Specifically, rather than directly maximize \min(R_1, R_2, \ldots, R_N), we’ll instead fix a parameter X and attempt to check whether it’s possible to make each R_i be \ge X.
If this is possible, the answer is surely at least X, otherwise it’s strictly less than X.
If it’s possible to make the answer \ge X, it’s also obviously possible to make it \ge X-1; so as long as we’re able to perform this check quickly enough, we can throw binary search on top to obtain the answer for an extra logarithmic factor.
That leaves us with the subproblem of: given X, can we make every R_i be \ge X by fixing a single edge?
To solve this, observe that as noted above, fixing a pipe (i, p_i) will affect the values of only all vertices in the subtree of i (and can never decrease these values - they’ll either remain the same or increase).
In other words, if R_i \lt X initially, we must change some edge that’s on the path from 1 to i.
More generally, if i_1, i_2, \ldots, i_m are all the vertices with R_{i_j} \lt X, the edge we change must lie on the path from 1 to all of these.
That’s possible if and only if the changed edge lies on the path from 1 to \text{LCA}(i_1, i_2, \ldots, i_m).
Here, \text{LCA} denotes the lowest common ancestor function.
So, let’s find all the vertices i_1, i_2, \ldots, i_m, and then from them compute L = \text{LCA}(i_1, \ldots, i_m).
We now need to decide which edge on the 1\to L path to fix.
Of course, trying them all is still not viable.
However, observe that we don’t actually care about which edge gets fixed - what really matters is the largest amount of water that can reach vertex L after fixing (at most) one edge on the 1\to L path.
This is because increasing the water that reaches L will naturally increase (or rather, not decrease) the water received by each of its descendants; so our best hope is to just maximize this value.
Finding this maximum is in fact quite easy, since it only depends on the 1\to L path now.
Let’s define dp_i to be the maximum amount of water that can reach vertex i if we’ve fixed a pipe already.
Then, we have
That is, we can either choose to fix the edge (i, p_i), in which case no previous edge would’ve been fixed (so p_i would’ve received R_{p_i} units of water, and i just receives all of them), or we would’ve fixed a previous edge, and by definition dp_{p_i} is the best that can reach p_i after which we compute the loss.
Armed with this dp, we know that dp_L is the maximum amount of water that can reach L after a single pipe fix on 1\to L.
Once this updated amount of water is known, simply simulate the process for the subtree of L, and check if all of them end up with values \ge X or not; and update the bounds of the binary search appropriately.
One thing that bears mention, is the computation of L = \text{LCA}(i_1, \ldots, i_m).
This can be done in two ways:
- The standard way, which is to note that \text{LCA}(i_1, \ldots, i_m) = \text{LCA}(i_1, \text{LCA}(i_2, \ldots, i_m)).
Using this, you can use any method of finding LCA quickly (binary lifting/RMQ/etc.) to compute the LCA. - Alternately, we can use the problem’s constraints to implement something much simpler, requiring no data structures at all.
Since p_i \lt i, L will just be the largest value that contains all of i_1, \ldots, i_m in its subtree.
This allows us to mark i_1, \ldots, i_m with the value 1, everything else with 0, and then find the largest vertex with a subtree sum of m - and so only a single DFS is needed.
Further, the p_i \lt i constraint allows for simple loops to substitute for a DFS to ease implementation further - the editorialist’s code below will serve as an example of how.
TIME COMPLEXITY:
\mathcal{O}(N\log K) per testcase.
CODE:
Editorialist's code (C++)
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int n, k; cin >> n >> k;
vector p(n, 0), a(n, 0);
for (int i = 1; i < n; ++i) {
cin >> p[i]; --p[i];
}
for (int i = 1; i < n; ++i) {
cin >> a[i];
}
vector<array<int, 2>> dp(n); // dp[i][j] = maximum water that can reach i after j pipe fixes
dp[0] = {k, k};
for (int i = 1; i < n; ++i) {
int x = a[i];
dp[i][0] = dp[p[i]][0] - (dp[p[i]][0] / x);
dp[i][1] = dp[p[i]][1] - (dp[p[i]][1] / x);
dp[i][1] = max(dp[i][1], dp[p[i]][0]);
}
vector ct(n, 0), nval(n, 0);
auto check = [&] (int x) {
// can only affect an edge that has all the values < x in its subtree
for (int &it : ct) it = 0;
for (int i = n-1; i >= 0; --i) {
ct[i] += dp[i][0] < x;
if (i) ct[p[i]] += ct[i];
}
int mx = ct[0];
if (mx == 0) return true;
int who = 0;
for (int i = 0; i < n; ++i) {
if (ct[i] == mx) who = i;
}
nval[0] = k;
for (int i = 1; i < n; ++i) {
if (i == who) nval[i] = dp[i][1];
else nval[i] = nval[p[i]] - (nval[p[i]] / a[i]);
if (nval[i] < x) return false;
}
return true;
};
int lo = 0, hi = k+1;
while (lo < hi) {
int mid = (lo + hi + 1) / 2;
if (check(mid)) lo = mid;
else hi = mid - 1;
}
cout << lo << '\n';
}
}