PermuTree (easy version)
Contest: Codeforces Round 890 (Div. 2) supported by Constructor Institute
Difficulty rating: 1800
Tags: dfs and similar, dp, trees

This is the easy version of the problem. The differences between the two versions are the constraint on $$$n$$$ and the time limit. You can make hacks only if both versions of the problem are solved.

You are given a tree with $$$n$$$ vertices rooted at vertex $$$1$$$.

For some permutation$$$^\dagger$$$ $$$a$$$ of length $$$n$$$, let $$$f(a)$$$ be the number of pairs of vertices $$$(u, v)$$$ such that $$$a_u < a_{\operatorname{lca}(u, v)} < a_v$$$. Here, $$$\operatorname{lca}(u,v)$$$ denotes the lowest common ancestor of vertices $$$u$$$ and $$$v$$$.

Find the maximum possible value of $$$f(a)$$$ over all permutations $$$a$$$ of length $$$n$$$.

$$$^\dagger$$$ A permutation of length $$$n$$$ is an array consisting of $$$n$$$ distinct integers from $$$1$$$ to $$$n$$$ in arbitrary order. For example, $$$[2,3,1,5,4]$$$ is a permutation, but $$$[1,2,2]$$$ is not a permutation ($$$2$$$ appears twice in the array), and $$$[1,3,4]$$$ is also not a permutation ($$$n=3$$$ but there is $$$4$$$ in the array).

Input format:
The first line contains a single integer $$$n$$$ ($$$2 \le n \le 5000$$$).

The second line contains $$$n - 1$$$ integers $$$p_2,p_3,\ldots,p_n$$$ ($$$1 \le p_i < i$$$) indicating that there is an edge between vertices $$$i$$$ and $$$p_i$$$.

Output format:
Output the maximum value of $$$f(a)$$$.

Constraints:
Time limit: 2.0 s
Memory limit: 512.0 MB

Note:
The tree in the first test:

One possible optimal permutation $$$a$$$ is $$$[2, 1, 4, 5, 3]$$$ with $$$4$$$ suitable pairs of vertices:

- $$$(2, 3)$$$, since $$$\operatorname{lca}(2, 3) = 1$$$ and $$$1 < 2 < 4$$$,
- $$$(2, 4)$$$, since $$$\operatorname{lca}(2, 4) = 1$$$ and $$$1 < 2 < 5$$$,
- $$$(2, 5)$$$, since $$$\operatorname{lca}(2, 5) = 1$$$ and $$$1 < 2 < 3$$$,
- $$$(5, 4)$$$, since $$$\operatorname{lca}(5, 4) = 3$$$ and $$$3 < 4 < 5$$$.

The tree in the third test:

The tree in the fourth test:

Example input:
5
1 1 3 3

Example output:
4

Example input:
2
1

Example output:
0

Editorial:
Let's consider the subproblem of maximizing the number of suitable pairs for some fixed $$$\operatorname{lca}(u, v) = x$$$. Then, we want to maximize the number of pairs $$$(u, v)$$$ such that $$$a_u < a_x < a_v$$$ and $$$u$$$ and $$$v$$$ are in different subtrees of $$$x$$$. So for each subtree of $$$x$$$, we only care about the number of vertices with $$$a_v > a_{x}$$$ and about the number of vertices with $$$a_v < a_x$$$.
Suppose vertex $$$x$$$ has $$$m$$$ subtrees and the $$$i$$$-th of them has $$$s_i$$$ vertices in it and $$$b_i$$$ vertices less with $$$a_v < a_x$$$. We will prove later that a suitable permutation $$$a$$$ exists for all possible correct values of $$$b_i$$$ (that is $$$0 \le b_i \le s_i$$$).
Then, the total number of suitable pairs (and the value we are trying to maximize) is equal to: $$$(s_1 - b_1) \cdot(0 + b_2 + \ldots + b_m) + (s_2 - b_2) \cdot(b_1 + 0 + \ldots + b_m) + \ldots + (s_m - b_m) \cdot(b_1 + b_2 + \ldots + 0)$$$.
Let $$$dp[i][B]$$$ be the maximum possible number of suitable pairs such that $$$u$$$ and $$$v$$$ lie in the first $$$i$$$ subtrees and $$$b_1 + b_2 + \cdots + b_i = B$$$ ($$$0 \le b_j \le s_j$$$ for all $$$1 \le j \le i$$$). Also define $$$S_i := s_1 + s_2 + \ldots + s_i$$$. Then:
• $$$dp[1][B] := 0$$$ for $$$0 \le B \le s_1$$$.
• $$$dp[i][B] := \max\limits_{\max(0, B - S_{i-1}) \le b_i \le \min(s_i, B)}{(dp[i - 1][B - b_i] + b_i \cdot (S_{i-1} - (B - b_i)) + (s_i - b_i) \cdot (B - b_i))}$$$ for $$$2 \le i \le m$$$ and $$$0 \le B \le S_{i}$$$.
In the statement above:
• The limits for $$$b_i$$$ are $$$\max(0, B - S_{i-1}) \le b_i \le \min(s_i, B)$$$ because:   $$$0 \le b_i \le s_i$$$ must hold,  $$$0 \le B - b_i \le S_{i - 1}$$$ must hold.
• $$$0 \le b_i \le s_i$$$ must hold,
• $$$0 \le B - b_i \le S_{i - 1}$$$ must hold.
• $$$dp[i - 1][B - b_i]$$$ is the number of suitable pairs $$$(u, v)$$$ in the first $$$i - 1$$$ subtrees.
• $$$b_i \cdot (S_{i-1} - (B - b_i))$$$ is the number of suitable pairs $$$(u, v)$$$ where $$$u$$$ lies in the $$$i$$$-th subtree and $$$v$$$ lies in the first $$$i - 1$$$ subtrees.
• $$$(s_i - b_i) \cdot (B - b_i)$$$ is the number of suitable pairs $$$(u, v)$$$ where $$$v$$$ lies in the $$$i$$$-th subtree and $$$u$$$ lies in the first $$$i - 1$$$ subtrees.
The maximum possible number of suitable pairs is equal to $$$\max\limits_{0 \le B \le S_m}{dp[m][B]}$$$.
Let's calculate the complexity of this dynamic programming subproblem. To do that, it's easier to consider iterating over $$$b_i$$$ and then iterating over all suitable values of $$$B$$$. Then, from the condition $$$0 \le b_i \le s_i$$$ we get $$$s_i + 1$$$ values of $$$b_i$$$, and from the condition $$$0 \le B - b_i \le S_{i - 1} \implies b_i \le B \le b_i + S_{i - 1}$$$ we get $$$S_{i - 1} + 1$$$ suitable values of $$$B$$$ for each $$$b_i$$$.
Taking into account that the $$$dp[i][B]$$$ array has $$$m \cdot (S_m + 1)$$$ values, adding everything up, the number of "operations" we get is equal to:
$$$$$$m \cdot (S_m + 1) + (s_2 + 1) \cdot (S_1 + 1) + (s_3 + 1) \cdot (S_2 + 1) + \ldots + (s_m + 1) \cdot (S_{m - 1} + 1) = $$$$$$
$$$$$$m \cdot (S_m + 1) + (s_2 \cdot S_1 + s_3 \cdot S_2 + \ldots + s_m \cdot S_{m-1}) + (s2 + \ldots + s_m) + (S_1 + \ldots + S_{m-1}) + (m - 1) \le $$$$$$
$$$$$$m \cdot S_m + m + (s_2 \cdot S_1 + s_3 \cdot S_2 + \ldots + s_m \cdot S_{m-1}) + S_m + (m - 1) \cdot S_m + m = $$$$$$
$$$$$$2m \cdot (S_m + 1) + (s_2 \cdot S_1 + s_3 \cdot S_2 + \ldots + s_m \cdot S_{m-1}) \le $$$$$$
$$$$$$2mn + (s_2 \cdot S_1 + s_3 \cdot S_2 + \ldots + s_m \cdot S_{m-1})$$$$$$
Returning to the original problem, let's solve the subproblem for each $$$x$$$ from $$$1$$$ to $$$n$$$ and add up the results. This is the upper bound on the answer, and we will prove a bit later that a suitable permutation $$$a$$$ exists.
Let's change our notation a bit:
• $$$m_x$$$ is the number of subtrees of vertex $$$x$$$.
• $$$s_{x,i}$$$ is the size of the $$$i$$$-th subtree of vertex $$$x$$$.
• $$$S_{x,i}$$$ is equal to $$$s_{x,1} + s_{x,2} + \ldots + s_{x,i}$$$.
• $$$b_{x,i}$$$ is the number of vertices in the $$$i$$$-th subtree for which $$$a_v < a_x$$$.
Adding up the number of "operations" for all subproblems, we get:
$$$$$$\displaystyle\sum_{x=1} ^{n} (2m_xn + s_{x,2} \cdot S_{x,1} + s_{x,3} \cdot S_{x,2} + \ldots + s_{x,m_x} \cdot S_{x,m_x-1}) = $$$$$$ $$$$$$2n \cdot (n - 1) + \displaystyle\sum_{x=1} ^{n} (s_{x,2} \cdot S_{x,1} + s_{x,3} \cdot S_{x,2} + \ldots + s_{x,m_x} \cdot S_{x,m_x-1})$$$$$$
Consider an undirected graph with $$$n$$$ vertices where there are initially no edges. You can think of the value of $$$s_{x,j} \cdot S_{x,j-1}$$$ as the number of edges added between each vertex of the first $$$j-1$$$ subtrees of $$$x$$$ and the $$$j$$$-th subtree of $$$x$$$. Since each pair of $$$(x, j)$$$ accounts for a unique set of edges of size $$$s_{x,j} \cdot S_{x,j-1}$$$, the sum $$$\sum_{x=1} ^{n} (\ldots)$$$ can be bounded above by $$$\frac{n \cdot (n - 1)}{2}$$$, the maximum number of unique edges in a graph with $$$n$$$ vertices:
$$$$$$2n \cdot (n - 1) + \displaystyle\sum_{x=1} ^{n} (s_{x,2} \cdot S_{x,1} + s_{x,3} \cdot S_{x,2} + \ldots + s_{x,m_x} \cdot S_{x,m_x-1}) \le$$$$$$
$$$$$$2n \cdot (n - 1) + \frac{n \cdot (n - 1)}{2}$$$$$$
So we get a complexity of $$$\mathcal{O}(n^2)$$$.
You can read more about this complexity analysis in section 7 of this Codeforces blog.
$$$ $$$
We solve independently for each value of $$$\operatorname{lca}(u,v)$$$ and add up the answers to get an upper bound for $$$f(a)$$$. We will prove by constructing $$$a$$$ that this bound is achievable.
Let $$$g(x, L)$$$ be a function that takes a vertex $$$x$$$ and a set of integers $$$L$$$, assigns some value from $$$L$$$ to $$$a_x$$$, and then calls $$$g(y, L_y)$$$ for all children of $$$x$$$ (direct descendants), where the sets $$$L_y$$$ do not intersect and their union is equal to $$$L \setminus \{ a_x \}$$$.
How it works:
1. Let $$$y_{1}, \ldots y_{m_x}$$$ be the children of vertex $$$x$$$ and $$$L_{y_i}$$$ initially be empty sets. Suppose that for vertex $$$x$$$ it is optimal to have the values of $$$b_{x,1}, b_{x,2}, \ldots, b_{x,m_x}$$$ be equal to $$$c_1, c_2, \ldots, c_{m_x}$$$ respectively.
2. For each $$$i$$$ from $$$1$$$ to $$$m_x$$$, take the $$$c_i$$$ smallest values from $$$L$$$, remove them from $$$L$$$, and add them to the set $$$L_{y_i}$$$.
3. Set $$$a_x$$$ to the smallest value in $$$L$$$ and remove it from the set.
4. For each $$$i$$$ from $$$1$$$ to $$$m_x$$$, take the $$$s_i - c_i$$$ smallest values from $$$L$$$, remove them from $$$L$$$, and add them to the set $$$L_{y_i}$$$.
5. For each $$$i$$$ from $$$1$$$ to $$$m_x$$$, call $$$g(y_i, L_{y_i})$$$.
Calling $$$g(x, \{1, 2, \ldots n \})$$$ constructs an optimal permutation $$$a$$$ for the given tree.
Complexity: $$$\mathcal{O}(n^2)$$$
Note: you can change $$$dp[i][B]$$$ to $$$dp[B]$$$ and iterate $$$B$$$ from $$$S_i$$$ to $$$0$$$. This improves the memory usage, but does not change the time complexity.