PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: satyam_343
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
DFS, Sorting
PROBLEM:
You’re given a tree with N vertices. Vertex i has value A_i.
The score of a tree is defined to be the minimum number of following operations required to make every A_i equal to 0:
- Choose a subset S of vertices of the tree such that for any x, y\in S such that x\neq y, x is not an ancestor of y.
- Then, reduce A_x by 1 for each x \in S.
Answer Q queries.
In each query, you’re given a value K. You want the score of the tree to be \leq K.
Find the minimum number of operations needed to achieve this, where in one operation you can choose two nodes u, v such that A_u \gt 0, then reduce A_u by 1 and increase A_v by 1.
In this version of the problem, Q = 1.
EXPLANATION:
Before worrying about queries and making the score small and whatnot, let’s first understand the setup itself properly. This means the first thing we do should be to figure out how to compute the score of a fixed tree with vertex values.
Consider some vertex u.
Observe that in one move, we can subtract at most 1 from a vertex that’s an ancestor of u (including u itself).
So, if S_u denotes the sum of values from the root to u, we definitely need \geq S_u moves to reduce everything to zero.
However, this analysis applies to every vertex u. This means we obtain \max(S) as a lower bound for the number of operations necessary.
It’s not hard to see that exactly \max(S) operations is also always achievable: one simple construction is to first perform A_1 operations at the root (which can’t include any other index anyway, since 1 is an ancestor of them all), then recursively solve for the children subtrees independently.
Since the children subtrees are unrelated to each other (in terms of ancestor relations), the number of operations will be the maximum that any of them need.
Repeating this shows that eventually we’ll use exactly \max(S) operations, as claimed.
Now, let’s think about answering queries.
This is the easy version of the task, so we only need to worry about answering a single one.
Our goal is to make the score \leq K.
As noted earlier, the score of a tree is just the maximum root → node path, so we want that to be \leq K.
Since all values are non-negative, the maximum root → node path will always be achieved at a leaf.
Hence, our objective is essentially to make sure that for every leaf v, S_v \leq K holds.
First, let’s get an edge case out of the way: note that the operation doesn’t change the sum of all values.
So, if there are L leaves in the tree, and L\cdot K \lt \text{sum}(A), it’s definitely impossible to make all the leaves have values \leq K, and the answer is -1.
If L\cdot K \geq \text{sum}(A), there’s enough “space” in the leaves: the answer will never be -1 (for example, by concentrating all values of A in the leaves, we can definitely ensure they’re all \leq K), so we only need to focus on minimizing the number of operations.
Each operation allows us to essentially move a value of 1 from one vertex to another.
It’s easy to see that the addition part of the operation should always be done to a leaf - if we add 1 to A_u, and u has a child c, adding 1 to A_c instead won’t increase the sum-on-path values till any leaf (and will in fact decrease some of them, if u has children other than c), which is why it’s always optimal to do so.
Similarly, the subtraction part should be done from “low depth” nodes, so as to reduce the values of as many leaves as possible.
Specifically, if we choose to subtract from A_u, then every ancestor of A_u should have a value of 0 already: otherwise it’s better to subtract from an ancestor.
Let’s split the process up into two “phases”: the subtraction phase and the addition phase.
In the subtraction phase, we’ll first perform only subtractions, from various nodes.
This will give us some extra space in the leaves (since some of them might have their S_u values reduced by the subtractions).
Then, we’ll perform the addition operations, where we try to fit all necessary additions into the space remaining in the leaves.
First, let’s look at the subtraction phase.
As noted earlier, it’s ideal to perform this starting from the root and moving downwards.
This means the process will look something like this:
- First, perform several subtractions at the root.
Here, our goal is to make S_u \leq K for every u.
So, we definitely need to perform \max(S_u) - K subtractions. On the other hand, we also can’t perform more than A_1 subtractions at the root.
This means we must perform \min(A_1, \max(S_u) - K) subtractions at the root (or 0, if this is negative). - Next, suppose we did indeed perform A_1 subtractions at the root.
The next step would be to perform subtractions on some of its children.
However, the children are independent of each other (in terms of root → node paths), so we can simply solve for each of them independently and add up their answers. - So, let’s look at a single child of the root, say c.
The situation is pretty similar here: at most A_c subtractions can be performed, but we also need to perform enough that every root → node path within this subtree becomes \leq K.
To this end, let’s define \text{sub}_c to be the maximum value of S_u across all u in the subtree of c.
Note that A_1 subtractions have been performed already, so the current maximum value is in fact \text{sub}_c - A_1. This means we need \min(A_c, \text{sub}_c - A_1 - K) operations here. - The same then applies to every node we visit, really: if we need to perform subtractions here, it means everything above it has been subtracted already.
So, for the node v, we need to perform \min(A_v, \text{sub}_v - (S_v - A_v) - K) operations (or 0, if this quantity is negative).
This completes the subtraction phase: we’ve subtracted just enough that every leaf has its sum-to-root value be \leq K.
Let X be the total number of subtractions done.
X is certainly a lower bound on the answer, since we’ve only performed subtractions that are necessary.
We now move to the addition phase, where we need to distribute X additions to the leaves while ensuring they remain \leq K.
First, let’s compute how much “free space” the leaves currently have among them.
This is easy to do: we know how many subtractions were done at each node, so just recompute root → leaf sums using this information.
Let S'_u denote the new root → u sum.
Then, for each leaf u, it has K - S'_u free space with which to perform additions.
Let F denote the sum of K - S'_u across all leaves u.
Then, if F \geq X, there is enough free space to perform the additions while ensuring that the leaves are all \leq K, so the answer is just X itself.
However, what if F \lt X?
In this case, we need more space in the leaves to accommodate all the additions.
The only way to obtain more space is by performing even more subtraction operations, however.
In particular, suppose we subtract 1 from A_y, for some vertex y.
Then, observe that X (the number of subtractions) increases by 1, while F (the available free space) increases by the number of leaves present in the subtree of y (which will henceforth be denoted L_y).
Effectively, the value of F - X (which we want to be \geq 0) has increased by L_y - 1.
Now, we’ve initially already performed all necessary subtractions to ensure that leaf scores are \leq K.
So, any additional subtractions won’t change this fact. This means we’re free to subtract from whichever nodes we want now - and since the goal is to make (F-X) \geq 0, it’s clearly best to choose whichever y has the largest L_y.
This immediately gives us a method of finding the number of extra operations needed.
Let’s sort all the vertices in descending order of the number of leaves in their subtrees.
Then, processing y in this order, perform as many subtractions as necessary (or possible) at y to ensure F-X becomes non-negative, where each subtraction increases this difference by L_y - 1.
Putting everything above together gives us a complete solution.
- Initially, for each vertex u, compute the following quantities with a DFS:
- S_u, the sum of values on the 1\to u path.
- \text{sub}_u, which is the maximum of S_v across all v in the subtree of u.
- L_u, the number of leaves in the subtree of u.
- If L_1\cdot K \lt \text{sum}(A), the answer is -1.
- First, the subtraction phase.
Define X to be the sum of \max(0, \min(A_u, \text{sub}_u - (S_u - A_u) - K)) across all u.
This is the number of subtractions necessary to make all root → node paths have a sum that’s \leq K. - Next, the addition phase.
Compute F to be the amount of free space in the leaves, after subtractions.- If F \geq X, the answer is X.
- Otherwise, process vertices y in descending order of their L_y values.
- For each y, perform as many subtractions as needed to make (F-X) become non-negative. Each subtraction increases the difference by L_y - 1 so this is some simple math.
Of course, A_y cannot fall below 0.
The time complexity is dominated by sorting vertices based on their leaf counts, so this is \mathcal{O}(N\log N) overall.
TIME COMPLEXITY:
\mathcal{O}(N\log N) 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, q; cin >> n >> q;
vector a(n, 0ll);
for (ll &x : a) cin >> x;
vector adj(n, vector<int>());
for (int i = 0; i < n-1; ++i) {
int u, v; cin >> u >> v;
adj[--u].push_back(--v);
adj[v].push_back(u);
}
ll totsum = accumulate(begin(a), end(a), 0ll), leafsum = 0;
vector<ll> sm(n), mx(n), leaves(n);
auto dfs = [&] (const auto &self, int u, int p) -> void {
bool leaf = true;
for (int v : adj[u]) if (v != p) {
sm[v] = sm[u] + a[v];
self(self, v, u);
mx[u] = max(mx[u], mx[v]);
leaves[u] += leaves[v];
leaf = false;
}
mx[u] = max(mx[u], sm[u]);
leaves[u] += leaf;
leafsum += sm[u]*leaf;
};
sm[0] = a[0];
dfs(dfs, 0, 0);
while (q--) {
ll k; cin >> k;
if (k*leaves[0] < totsum) {
cout << -1 << '\n';
continue;
}
vector<array<ll, 2>> extra;
ll ops = 0, space = 0;
for (int i = 0; i < n; ++i) {
ll cur = max(0ll, min(1ll*a[i], mx[i] - (sm[i] - a[i]) - k));
extra.push_back({leaves[i], a[i] - cur});
ops += cur;
space += cur*leaves[i];
}
space = k*leaves[0] - (leafsum - space);
ranges::sort(extra);
ranges::reverse(extra);
ll rem = max(0ll, ops - space);
for (auto [val, ct] : extra) {
// ct operations, each can reduce by (val - 1)
// min(ct, ceil(rem / (val - 1))) ops
if (val == 1) break;;
ll take = min(1ll*ct, (rem + val - 2) / (val - 1));
ops += take;
rem = max(0ll, rem - (val-1)*take);
}
if (rem > 0) cout << -1 << '\n';
else cout << ops << '\n';
}
}
}