PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
DFS, Dynamic Programming, Sorting
PROBLEM:
You’re given a tree on N vertices.
For a permutation P of \{1, \ldots, N\}, we define \text{score}(P) as follows:
- For vertex u, let M_u denote the median of all P_v such that \text{dist}(u, v) \le 1.
- u is called good if M_u = P_u.
- \text{score}(P) is the number of good vertices in the tree.
Find the maximum possible value of \text{score}(P) across all choices of P.
EXPLANATION:
Consider some vertex u. Suppose we want u to be the median of its and its neighboring values.
Then,
- If u has an even number of neighbors - say 2k - then k of them must have a smaller value than P_u and the other k must have a larger value.
- If u has an odd number of neighbors - say 2k+1 - then k of them must have a smaller value and k+1 of them must have a larger value than P_u.
In particular, note that we only care about relative values here - that is, for each edge (u, v), we only care about whether P_u \gt P_v or not.
With this line of thought, we can think of giving each edge an orientation.
That is, for the edge (u, v), we can direct it u\to v if we want P_u \gt P_v, and direct it u\gets v otherwise.
Suppose we make this choice independently for each of the edges of the tree (so there are 2^{N-1} possible orientations).
Observe that we can always find a permutation that satisfies the orientation of edges: a simple construction is to take any topological sort of the resulting directed graph (which will always exist, since it has no cycles), and fill values in descending order following the toposort.
So, rather than finding a permutation that maximizes \text{score}(P), we can instead work on finding an orientation that maximizes \text{score}(P).
This is helpful, because as noted above the orientation can be decided independently for each edge, making it a much more local quantity to work with - which is great since we only care about neighbors of any given vertex.
Consider some vertex u.
Let’s understand how we want to orient the edges incident to u.
There are two possibilities:
- We choose to not make u a good vertex, i.e. it doesn’t equal the median of its neighbors.
Here, observe that all edges incident to u are now entirely “free”, i.e. we can just decide their directions however we like.
Further, all the children of u are now entirely independent of each other (and the rest of the tree), so they can be solved for separately and the answers added up.
As for the parent of u, we can again orient the edge between them however we like. - We choose to make u a good vertex.
Here, we’re more constrained: half the edges incident to u must be pointing toward it, and the other half must be away from it.
Once again though, if we decide which half of the edges point toward u and which don’t, the children subtrees of u become independent problems, just with an additional restriction on their topmost edge.
This hints toward a solution using dynamic programming.
Specifically, for each u, let’s define:
- dp_1(u) is the maximum possible score considering only the subtree of u, with the constraint that the edge between u and \text{parent}(u) is pointed towards u.
- dp_2(u) is the maximum possible score considering only the subtree of u, with the constraint that the edge between u and \text{parent}(u) is pointed towards \text{parent}(u).
We now have to work on transitions.
Let’s look at just dp_1(u) first.
As noted above, there are two possibilities: either we make u good, or we don’t.
If we don’t make u good, each child can have its orientation decided as we like.
Thus, for each child c of u, we can take either dp_1(c) or dp_2(c), whichever is better.
So, we obtain a value of \sum_{c} \max(dp_1(c), dp_2(c)) across all children c of u.
Next, we consider the case where u is made good.
We’re in the dp_1(u) case, meaning we have one edge pointing towards u already (that being \text{parent}(u) \to u).
That tells us exactly how many edges from the children of u must point towards u.
Let this necessary count be k.
We now want to choose the value dp_2(c) for exactly k children of u, and dp_1(c) for all the others; while maximizing the sum of the chosen values.
That can be done as follows:
- First, start out by taking just dp_1(c) for all c.
- Next, for each c, observe that we obtain a “profit” of (dp_2(c) - dp_1(c)) by taking dp_2 here instead.
- We want to switch exactly k children; clearly it’s optimal to just take the largest k profits among them.
This can be done easily by sorting the list of profits.
If the maximum profit here is P, the overall value we obtain is 1 + P + \sum_c dp_1(c).
Thus, we have dp_1(u) = \max(1 + P + \sum_c dp_1(c), \sum_{c} \max(dp_1(c), dp_2(c))).
dp_2(u) can be computed similarly (the only thing that will change is that k will change by 1).
So, for a node u, we’re able to compute dp_1(u) and dp_2(u) in \mathcal{O}(\text{deg}(u)\log \text{deg}(u)) time.
Summed up across all u, this is bounded by \mathcal{O}(N\log N) (because the sum of degrees is 2N-2), so it’s fast enough for us.
It’s also possible to optimize this to linear time using a selection algorithm instead of sorting, though that of course is not needed to get AC.
TIME COMPLEXITY:
\mathcal{O}(N\log N) per testcase.
CODE:
Editorialist's code (PyPy3)
for _ in range(int(input())):
n = int(input())
par = [0, 0] + list(map(int, input().split()))
dp = [[0, 0] for i in range(n+1)] # 0: u <- par[u], 1: u -> par[u]
children = [ [] for i in range(n+1)]
for i in range(2, n+1):
children[par[i]].append(i)
for i in reversed(range(1, n+1)):
sm, best = 0, 0
difs = []
for v in children[i]:
best += max(dp[v][0], dp[v][1])
sm += dp[v][0]
difs.append(dp[v][1] - dp[v][0])
half = (len(children[i]) + (1 if i > 1 else 0)) // 2
difs.sort(reverse=True)
# u -> par[u]
dp[i][1] = max(best, 1 + sm + sum(difs[:half]))
# u <- par[u]
if half > 0: dp[i][0] = max(best, 1 + sm + sum(difs[:half-1]))
else: dp[i][0] = best
print(dp[1][1])