PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: evenvalue
Tester: apoorv_me
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
DFS, dynamic programming
PROBLEM:
You have a tree with N vertices rooted at vertex 1, and an integer K.
Starting at vertex 1, you can do the following operations:
- If you’re standing on a leaf, consider it conquered.
Then, choose some k (1 \leq k \leq K), and move to the k-th ancestor of this leaf. - If you’re not on a leaf, move to one of its children.
What’s the maximum number of distinct leaves that can be conquered?
EXPLANATION:
The only way to move up the tree is to reach a leaf and then pick one of its ancestors.
This means an optimal sequence of moves will look as follows:
- Start at vertex 1, and choose one of its children (say v_1).
Move to v_1. - Perform some moves within the subtree of v_1, conquering as many leaves as possible.
Then, use a leaf to jump back to vertex 1, and never enter v_1 again. - Next, pick another child v_2, and do the same thing.
\vdots - Finally, pick a child v_k of 1, from which you won’t return to 1 (either by choice, or because it’s impossible to do so).
Move to v_k, and solve for it.
This leads us to a dynamic programming solution.
Define:
- \text{ans}[u] to be the maximum number of leaves that can be conquered in the subtree of u.
- \text{return}[u] to be the maximum number of leaves that can be conquered in the subtree of u, if you start at u and are still able to return to it afterwards.
The answer we want is \text{ans}[1].
Let’s attempt to compute these two arrays, starting with \text{return}.
Suppose you’re at a vertex u, and v is one of its children.
If it’s possible to visit v, perform some moves in the subtree of v, and then return to u, we make the following observations:
- First, there should be a leaf that’s “high enough” that it can be used to return to u at all.
That is, there should exist a leaf x in the subtree of x such that the depth of x is at most \text{dep}[u] + K.
This is easy to check: it’s enough to compute the highest leaf in the subtree of v, and then check if it satisfies this condition. - Second, notice that it’s also possible to use the same sequence of moves to end up at v instead of u (on the final upward move, you can choose to end up at v rather than u).
Conversely (provided the “high enough” leaf exists), any sequence of moves that starts at v and ends up at v can be extended to end at u instead, by first ending at v, then moving down to this leaf and then moving up to u.
This means the best value that can be obtained from this subtree is exactly \text{return}[v].
This makes \text{return}[u] quite simple to compute.
For each child v of u, if there exists a high enough leaf in the subtree of v, add \text{return}[v] to \text{return}[u].
If u is a leaf, \text{return}[u] = 1.
Computing the highest leaves in every subtree is easily done with a DFS.
Now that we know all the \text{return}[u] values, let’s compute the \text{ans}[u] values.
As noted at the start, when computing \text{ans}[u], we can choose at most one child v, and go into it without ever returning to u.
That is, for every child v of u except maybe one, the value we take will still be just \text{return}[v] (or 0, if a high enough leaf doesn’t exist).
In other words, we start with \text{ans}[u] = \text{return}[u], and then for exactly one child, we’ll replace \text{return}[v] with \text{ans}[v].
Since we want to maximize the sum, clearly it’s best to choose whichever child has the highest value of \text{ans}[v] - \text{return}[v].
So, quite simply, all that needs to be done is to find the highest value of \text{ans}[v] - \text{return}[v] across all children; and set \text{ans}[u] to be \text{return}[u] plus this maximum.
This is easily done with a DFS, and can even just be done when computing the \text{return}[u] values.
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
namespace Read {
int Int() {
int x;
cin >> x;
return x;
}
}//namespace Read
constexpr int kInf = 1e9 + 10;
inline void solution() {
const int n = Read::Int();
const int k = Read::Int();
vector<vector<int>> g(n);
for (int i = 0; i < n - 1; i++) {
const int x = Read::Int() - 1;
const int y = Read::Int() - 1;
g[x].push_back(y);
g[y].push_back(x);
}
vector<int> depth(n);
vector<int> shallow(n, kInf);
auto is_sink = [&](const int x) {
return shallow[x] - depth[x] >= k;
};
vector<int> recoil(n);
vector<int> best(n);
function<void(int, int)> dfs = [&](const int x, const int p) {
if (p != -1 and g[x].size() == 1) {
shallow[x] = depth[x];
recoil[x] = 1;
best[x] = 1;
}
for (const int y : g[x]) {
if (y == p) continue;
depth[y] = depth[x] + 1;
dfs(y, x);
shallow[x] = min(shallow[x], shallow[y]);
if (not is_sink(y)) {
recoil[x] += recoil[y];
}
}
int extra = 0;
for (const int y : g[x]) {
if (y == p) continue;
extra = max(extra, best[y] - (is_sink(y) ? 0 : recoil[y]));
}
best[x] = recoil[x] + extra;
};
dfs(0, -1);
cout << best[0] << '\n';
}
int32_t main() {
ios_base::sync_with_stdio(false);
cin.tie(nullptr);
int testcases = 1;
cin >> testcases;
while (testcases--) {
solution();
}
}
Tester's code (C++)
#include <bits/stdc++.h>
using namespace std;
int main() {
cin.tie(0) -> sync_with_stdio(0);
int T; cin >> T;
while(T-- > 0) {
int N, K; cin >> N >> K;
vector<vector<int>> adj(N);
for(int i = 1 ; i < N ; ++i) {
int u, v; cin >> u >> v;
adj[u - 1].push_back(v - 1);
adj[v - 1].push_back(u - 1);
}
int ans = 1;
function<array<int, 3>(int, int)> dfs = [&](int node, int par) -> array<int, 3> {
vector<array<int, 3>> sol;
for(auto &u: adj[node]) if(u != par) {
sol.push_back(dfs(u, node));
};
if(sol.empty()) {
return {1, 1, 1};
}
sort(sol.rbegin(), sol.rend());
int count = 0, hh = sol.front()[0];
int mx = 0;
for(auto &[h, v, b]: sol) {
if(h <= K) {
count += v;
mx = max(mx, b - v);
} else {
mx = max(mx, b);
}
hh = h;
}
ans = max(ans, count + mx);
return {hh + 1, count, count + mx};
};
dfs(0, -1);
cout << ans << '\n';
}
}
Editorialist's code (C++)
// #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;
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int n, k; cin >> n >> k;
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);
}
vector<int> dp1(n), dp2(n);
auto dfs = [&] (const auto &self, int u, int p, int dep = 0) -> int {
int mnleaf = n+1;
bool isleaf = true;
int bestinc = 0;
for (int v : adj[u]) {
if (v == p) continue;
int high = self(self, v, u, dep + 1);
isleaf = false;
mnleaf = min(mnleaf, high);
if (dep + k >= high) {
dp2[u] += dp2[v];
bestinc = max(bestinc, dp1[v] - dp2[v]);
}
else {
bestinc = max(bestinc, dp1[v]);
}
}
if (isleaf) {
dp1[u] = dp2[u] = 1;
return dep;
}
dp1[u] = dp2[u] + bestinc;
return mnleaf;
};
dfs(dfs, 0, 0);
cout << dp1[0] << '\n';
}
}
