PermuTree (hard version)
Contest: Codeforces Round 890 (Div. 2) supported by Constructor Institute
Difficulty rating: 2700
Tags: bitmasks, dfs and similar, dp, fft, greedy, implementation, math, trees

This is the hard 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 10^6$$$).

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: 3.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:
First of all, the most important takeaways from the editorial for the easy version:
• For each $$$x$$$ from $$$1$$$ to $$$n$$$, you can solve $$$n$$$ independent subproblems of maximizing the number of suitable pairs for a fixed value of $$$\operatorname{lca}(u, v) = x$$$, and then add up the results to get the final answer.
• For each subproblem, the number of suitable pairs of vertices can be calculate based on the sizes of all the child subtrees $$$s_{x,i}$$$ and the number of vertices with $$$a_v < a_x$$$ in each of them, which we will call $$$b_{x,i}$$$.
Consider the subproblem for some $$$x$$$. Let $$$m_x$$$ be the number of child subtrees of vertex $$$x$$$. Suppose some optimal values of $$$b_{x,1}, b_{x,2}, \ldots b_{x,m_x}$$$ that maximize the number of suitable pairs are $$$c_1, c_2, \ldots c_{m_x}$$$ respectively.
$$$ $$$
Claim: there exists an optimal solution where $$$c_i = 0$$$ or $$$c_i = s_{x,i}$$$ for all $$$i$$$ from $$$1$$$ to $$$m_x$$$.
Proof: let's consider some optimal solution $$$c_1, c_2, \ldots c_{m_x}$$$ where $$$0 < c_i < s_{x,i}$$$ for some $$$i$$$. Suppose $$$D$$$ is the number of suitable pairs of vertices such that both $$$u$$$ and $$$v$$$ don't lie in the $$$i$$$-th child's subtree. Then, the total number of suitable pairs is equal to: $$$$$$D + c_i \cdot \sum_{j \neq i}{(s_{x, j} - c_j)} + (s_{x, i} - c_i) \cdot \sum_{j \neq i} c_j$$$$$$ If we define $$$C_1 := \sum_{j \neq i}{(s_{x, j} - c_j)}$$$ and $$$C_2 := \sum_{j \neq i} c_j$$$, we get an even simpler formula: $$$$$$D + c_i \cdot C_1 + (s_{x, i} - c_i) \cdot C_2 =$$$$$$ $$$$$$(D + C_2 \cdot s_{x,i}) + c_i \cdot (C_1 - C_2)$$$$$$
The values of $$$(D + C_2 \cdot s_{x,i})$$$, $$$C_1$$$, and $$$C_2$$$ don't depend on the value of $$$c_i$$$, so we can freely change it to anything from $$$[0; s_{x, i}]$$$ to maximize the value of the formula. So:
• If $$$C_1 \ge C_2$$$, the value of the formula will not decrease if we set $$$c_i$$$ to $$$s_{x,i}$$$.
• If $$$C_1 \le C_2$$$, the value of the formula will not decrease if we set $$$c_i$$$ to $$$0$$$.
Since at least one of this conditions has to be true, we can always take some optimal solution with $$$k$$$ indices such that $$$0 < c_i < s_{x, i}$$$ and get an optimal solution with $$$k - 1$$$ such indices. We repeat this process until $$$k$$$ becomes $$$0$$$, which means that $$$c_i = 0$$$ or $$$c_i = s_{x, i}$$$ for all $$$i$$$ from $$$1$$$ to $$$m_x$$$, as was stated in the claim.
$$$ $$$
Consider an optimal solution $$$c_1, c_2, \ldots c_{m_x}$$$ where $$$c_i$$$ is equal to either $$$0$$$ or $$$s_{i, x}$$$. Let $$$S_1$$$ be the sum of $$$s_{i, x}$$$ over all $$$i$$$ such that $$$c_i = 0$$$ and $$$S_2$$$ be the sum of $$$s_{i, x}$$$ over all $$$i$$$ such that $$$c_i = s_{x, i}$$$. Then, the number of suitable pairs is equal to: $$$$$$S_1 \cdot S_2$$$$$$
If we define $$$S$$$ as $$$s_{x, 1} + s_{x, 2} + \ldots s_{x,m_x}$$$, we can see that $$$S_2$$$ is equal to $$$S - S_1$$$, so we are actually maximizing the value of: $$$$$$S_1 \cdot (S - S_1)$$$$$$
To do this, we can iterate $$$S_1$$$ from $$$0$$$ to $$$S$$$, find if it is possible to choose some subset of indices $$$I \subseteq \{1, 2, \ldots, m_x \}$$$ such that $$$\displaystyle\sum_{i \in I}{s_{x,i}} = S_1$$$, and update the current best answer if such a set $$$I$$$ exists.
Doing this efficiently is known as the subset sum problem (or some variation of it). One of the simplest ways to solve it in this particular case with a $$$\mathcal{O}(m_x \cdot S)$$$ dynamic programming solution: $$$dp[j][S_1]$$$ is equal to $$$1$$$ if there exists a set of indices $$$I \subseteq \{ 1, 2, \ldots, j \}$$$ such that $$$\displaystyle\sum_{i \in I}{s_{x,i}} = S_1$$$ and is equal to $$$0$$$ otherwise. We will proceed to optimize it.
The first standard optimization will bring the complexity down to $$$\mathcal{O}(S \sqrt S)$$$, you can read about it in Section 3 of this Codeforces blog or in the "Subset Sum Speedup 1" section of this Codeforces blog.
$$$ $$$
Since we are trying to maximize the value of $$$S_1 \cdot (S - S_1)$$$, the optimal choice will also minimize the value of $$$| (S - S_1) - S_1 | = | S - 2S_1 |$$$.
Claim: denote $$$\max(s_{x,1}, \ldots, s_{x,m_x})$$$ as $$$M_x$$$. If $$$M_x \ge \frac{S}{2}$$$, the smallest possible value of $$$| S - 2S_1 |$$$ is achieved with $$$S_1 = M_x$$$ or $$$S_1 = S - M_x$$$.
Proof: it is clearly possible to get $$$S_1 = M_x$$$. Suppose it is possible to get some value of $$$S_1$$$ in the segment $$$[S - M_x + 1; M_x - 1]$$$. Then, $$$M_x$$$ did not contribute to the sum $$$S_1$$$, and so $$$S_1 + M_x$$$ can't exceed $$$S$$$. But $$$S_1 + M_x \ge S - M_x + 1 + M_x = S + 1$$$, so we get a contradiction, hence it is not possible to get some $$$S_1 \in [S - M_x + 1; M_x - 1]$$$.
So if we have $$$M_x \ge \frac{S}{2}$$$, we actually don't have to run the dynamic programming part. We will now analyze the complexity of the solution with this optimization.
$$$ $$$
Let's now define $$$S_x$$$ as $$$s_{x, 1} + \ldots + s_{x, m_x}$$$. Let's also call a vertex $$$x$$$ heavy if $$$M_x \ge \frac{S_x}{2}$$$ and light otherwise.
Suppose that vertex $$$x$$$ has $$$k$$$ light ancestors on it's path to the root:
• Every subtree of a light vertex $$$v$$$ has at most $$$\frac{S_v}{2}$$$ vertices in it.
• Every subtree of a heavy vertex $$$v$$$ has at most $$$S_v - 1$$$ vertices in it.
For each $$$k$$$ from $$$0$$$ to $$$n$$$, let $$$L_k$$$ be the set of light vertices with exactly $$$k$$$ light ancestors (a vertex doesn't count as it's own ancestor):
• Note that if two light vertices $$$u$$$ and $$$v$$$ are in the same set $$$L_k$$$, one can't be the ancestor of the other. So, the sum $$$\sum_{x \in L_{k}} S_x$$$ does not exceed $$$n$$$.
• $$$S_x \le \frac{n}{2^k}$$$ for all $$$x \in L_{k}$$$.
For some $$$k$$$, let's calculate the total number of "operations" taken in the dynamic programming part for all $$$x \in L_k$$$. To do that, we:
• Iterate over the number of vertices in the set $$$L_{k}$$$, let's call it $$$p$$$, from $$$0$$$ to $$$n$$$.
• Let $$$d_1, d_2, \ldots, d_p$$$ be the sizes of the subtrees of vertices in $$$L_{k}$$$. Then we need to maximize the value of: $$$$$$d_1 \sqrt d_1 + d_2 \sqrt d_2 + \ldots + d_p \sqrt d_p$$$$$$ over all $$$d_1, d_2, \ldots, d_p$$$ such that $$$1 \le d_n \le \frac{n}{2^k}$$$ and $$$d_1 + \ldots + d_n \le n$$$.
You might have noticed that the above process can be simplified by allowing $$$d_i$$$ to be equal to $$$0$$$. In that case, the number of vertices $$$p$$$ is equal to the number of positive integers in $$$d_1, d_2, \ldots, d_n$$$, and we need to maximize:
$$$$$$d_1 \sqrt d_1 + d_2 \sqrt d_2 + \ldots + d_n \sqrt d_n$$$$$$
over all integer $$$d_1, d_2, \ldots , d_n$$$ such that $$$0 \le d_n \le \frac{n}{2^k}$$$ and $$$d_1 + \ldots + d_n \le n$$$.
Let $$$t := \left\lfloor \frac{n}{2^k} \right\rfloor$$$. We'll assume $$$t \ge 1$$$, otherwise, $$$d_i$$$ is equal to $$$0$$$ for all $$$1 \le i \le n$$$ and the sum is always equal to $$$0$$$.
Lemma 1: it is optimal to have $$$d_1 + d_2 + \ldots + d_n = n$$$.
Proof: if we can increase some $$$d_i$$$, it is clearly optimal to do so.
Lemma 2: suppose the maximum is achieved on the set of values $$$e_1, e_2, \ldots, e_n$$$. Then there does not exist a pair of indices $$$i$$$ and $$$j$$$ such that $$$0 < e_i \le e_j < t$$$.
Proof: define $$$f(x) := \left( x + 1 \right) \sqrt {x + 1} - x \sqrt x$$$. Note that $$$f(x)$$$ is a monotonically increasing function for all $$$x \ge 0$$$. You can prove it by, for example, analyzing it's derivative: $$$\frac{3}{2} (\sqrt {x + 1} - \sqrt x)$$$.
If we decrease $$$e_i$$$ by $$$1$$$ and increase $$$e_j$$$ by $$$1$$$, the sum changes by: $$$$$$(e_j + 1) \sqrt {e_j + 1} + (e_i - 1) \sqrt {e_i - 1} - e_j \sqrt e_j - e_i \sqrt e_i =$$$$$$ $$$$$$ f(e_j) - f(e_i - 1) > 0$$$$$$
By doing this, we haven't broken any constraints on the values of $$$e_i$$$, but have increased the total sum, so the previous values of $$$e$$$ couldn't have been optimal.
This lemma gives us the nice property that the optimal values of $$$e_1, e_2, \ldots, e_n$$$ look something like $$$t, \ldots, t, y, 0, \ldots, 0$$$; where $$$0 < y < t$$$.
With this we can conclude that the total number of "operations" performed for all vertices in $$$L_k$$$ is at most: $$$$$$\left\lceil \frac{n}{t} \right\rceil \cdot t \sqrt t \le \left( \frac{n}{t} + 1 \right) \cdot t \sqrt t \le (n + t) \sqrt t \le 2n \sqrt t$$$$$$
Adding this up over all $$$k$$$ with $$$t \ge 1$$$, we get: $$$$$$2n \cdot \left(\sqrt{\left\lfloor \frac{n}{1} \right\rfloor} + \sqrt{\left\lfloor \frac{n}{2} \right\rfloor} + \sqrt{\left\lfloor \frac{n}{4} \right\rfloor} + \sqrt{\left\lfloor \frac{n}{8} \right\rfloor} + \ldots + 1 \right) \le$$$$$$ $$$$$$2n \cdot \left( \sqrt{\frac{n}{1}} + \sqrt{\frac{n}{2}} + \sqrt{\frac{n}{4}} + \sqrt{\frac{n}{8}} + \ldots \right) =$$$$$$ $$$$$$2n \sqrt n \cdot \left(1 + \frac{1}{\sqrt{2}} + \frac{1}{\sqrt{4}} + \frac{1}{\sqrt{8}} + \ldots \right) $$$$$$
We can notice that the sum in the brackets in bounded by the sum of an infinite geometric progression with a sum of $$$\frac{\sqrt{2}}{\sqrt{2} - 1} \le 4$$$, so the total sum is: $$$$$$\le 2n \sqrt n \cdot 4 = 8n \sqrt n = \mathcal{O}(n \sqrt n)$$$$$$
Finally, we will also use bitsets to speed up our dynamic programming part and bring the complexity down to $$$\mathcal{O}(\frac{n \sqrt n}{w})$$$, where $$$w$$$ is the word size, usually $$$32$$$ or $$$64$$$.
You can learn about bitsets in this Codeforces blog. To get a simple variable length bitset in C++, refer to this comment.
Complexity: $$$\mathcal{O}(\frac{n \sqrt n}{w})$$$.