PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: sheercold and Codula
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Dynamic Programming, DFS
PROBLEM:
You have a tree with N vertices. Assign a permutation P of [0, N-1] to the vertices of the tree so that \displaystyle\sum_{u=1}^N\sum_{v=u}^N f(u, v) is maximized, where f(u, v) is the mex of labels on the path between u and v.
EXPLANATION:
Let’s try to analyze what an optimal configuration could possibly look like.
First, note that for a path to have a MEX of x, it must contain all the values 0, 1, 2, \ldots, x-1.
Since we’re working with a permutation of [0, N-1], there’s only one occurrence of each of these values, so there’s no worry about duplicates.
In particular, let C_x be the number of paths containing all the values 0, 1, 2, \ldots x.
C_x is essentially the number of paths whose MEX is at least x+1.
Then, note that the score of a tree is exactly C_0 + C_1 + C_2 + \ldots
This is useful to keep in mind.
Suppose we start with an arbitrary assignment of labels - let’s figure out how to improve this.
Let u_x denote the vertex to which label x is assigned.
First, look at u_0 and u_1.
It’s always optimal for these two to be adjacent to each other: if they aren’t, then if v is the first vertex on the u_0 \to u_1 path, we could just swap the labels of v and u_1 instead: compared to the older assignment, it can be seen that the newer assignment is strictly better, because:
- Any path that contained both 0 and 1 in the older assignment, still does.
In fact, any path that contained all of 0, 1, 2, \ldots, k for any k, still does.
Thus, no path had its MEX reduced. - However, some paths that contained only 0 earlier, now contain both 0 and 1.
So, their MEX strictly increased - which in combination with the first point means the overall sum of MEX-es increased.
Now that u_0 and u_1 are adjacent, let’s look at u_2.
Similar reasoning shows that this should be adjacent to either u_0 or u_1: if it isn’t, just move it towards one of them as far as it can go and the overall answer increases.
So, u_0, u_1, u_2 will all be on a single path of length 3 (in some order).
Next, we look at u_3.
Pretty much the exact same reasoning holds: if it’s not adjacent to the path formed by the previous values, then placing it adjacent to one endpoint of that path (and hence extending the path) will always be better.
In general, this tells us what the structure of an optimal solution will look like.
We start with some vertex having the label 0.
Then, some neighbor of this vertex receives the value 1, then some neighbor of one of the previous values receives 2, then some neighbor of the length-3 path receives label 3, some neighbor of the now length-4 path receives label 4, and so on.
Or, more generically: we start with the path being \{0\} at some vertex, and then repeatedly extend the path on some side by the next missing value, till it’s no longer possible to do so (i.e. when both endpoints of the path are leaves).
In particular, note that any any point of time, one endpoint of the path will contain the maximum value on it.
Our goal is now to figure out what choices are optimal here: where to start the path with \{0\} and then how to expand.
To do that, we’ll employ dynamic programming.
For two vertices u and v, define L(u, v) to be the number of vertices on the path between u and v.
Define dp(u, v) to be the maximum possible score if we’ve assigned values 0,1 , 2, \ldots, L(u, v)-1 to the path between u and v.
We now need to compute this DP.
To do that, we need to consider two cases.
First, suppose u \neq v.
Let f_u be the first vertex on the u \to v path (other than u), and f_v be the first vertex on the v\to u path (other than v).
Then, as noted above, either u or v must contain the largest value (which is L(u, v)-1).
Suppose we place L(u, v)-1 at u.
Then, the values 0, 1, 2, \ldots, L(u,v)-2 must appear in the remainder of the path excluding u, which is exactly the path between f_u and v.
By definition, the best score we can obtain out of this path is dp(f_u, v).
That only leaves the contribution of the entire path (u, v) itself.
Note that the path (u, v) will just add 1 to the score for every path that contains (u, v).
So, all we need to do is count how many such paths there are.
That’s reasonably simple to do: for example you can do casework on whether one of u or v is an ancestor of the other one or not (after rooting at some vertex).
If neither is an ancestor of the other, the number of relevant paths is the product of their subtree sizes; and if one is an ancestor of the other then the ancestor will have everything outside its subtree rather than everything inside.
So, if there are ct(u, v) paths that contain the (u, v) path, then we obtain an overall score of dp(f_u, v) + ct(u, v) by placing the largest value at u.
Similar analysis shows that placing the value at v instead results in a value of dp(u, f_v) + ct(u, v).
Thus, we just have
Note that the reason this works is because of the very first observation made: that the score of a tree equals the sum of “number of paths that contain 0, 1, 2, \ldots, x” across all values of x.
As noted above, ct(u, v) can be computed in constant time once subtree sizes and ancestor relations are known, so that’s fast.
Computing f_u and f_v is also a fairly standard task: once again you can case on ancestor-descendant relations to compute the appropriate vertex, or you can use the fact that N is small enough to just precompute the appropriate vertex in quadratic time.
In fact, you can even just brute-force compute the value, by iterating through all neighbors of u and checking which of them is closer to v - this is still quadratic time (if distances are precomputed).
Either way, computing dp(u, v) can be done quickly: all values can be computed in quadratic time.
This leaves the u=v case.
This is fairly simple: dp(u, u) simply equals the number of paths containing u; because only the value 0 is placed.
To compute this, once again you can use subtree sizes, since the only paths that don’t pass through u are those that lie entirely inside one of the smaller trees obtained when removing u from the tree.
Once all the dp(u, v) values are known, the answer is simply the largest of them.
TIME COMPLEXITY:
\mathcal{O}(N^2) 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; cin >> n;
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 first(n, vector(n, -1)); // first[u][v] = first vertex on u->v path
vector subsz(n, 0);
auto dfs = [&] (const auto &self, int u, int p, auto &cur) -> void {
subsz[u] = 1;
for (int v : adj[u]) if (v != p) {
if (p == -1) cur[v] = v;
else cur[v] = cur[u];
self(self, v, u, cur);
subsz[u] += subsz[v];
}
};
for (int i = 0; i < n; ++i)
dfs(dfs, i, -1, first[i]);
vector dp(n, vector(n, 0ll));
auto calc = [&] (const auto &self, int u, int v) -> ll {
if (u > v) swap(u, v);
if (dp[u][v]) return dp[u][v];
if (u == v) {
dp[u][u] = n*(n+1)/2;
for (int x : adj[u]) {
int s = subsz[x];
if (s > subsz[u]) s = n - subsz[u];
dp[u][u] -= s*(s+1)/2;
}
return dp[u][u];
}
auto &res = dp[u][v];
res = max(self(self, first[u][v], v), self(self, u, first[v][u]));
ll x = subsz[u], y = subsz[v];
if (subsz[first[u][v]] < subsz[u]) { // u is ancestor of v
x = n - subsz[first[u][v]];
}
if (subsz[first[v][u]] < subsz[v]) { // v is ancestor of u
y = n - subsz[first[v][u]];
}
res += x*y;
return res;
};
ll ans = 0;
for (int i = 0; i < n; ++i)
for (int j = i; j < n; ++j)
ans = max(ans, calc(calc, i, j));
cout << ans << '\n';
}
}