For Subtask 1, because all connection points satisfy the constraint ~x^2+x=y_i~ for some integer ~x~, the problem is equivalent to finding the diameter of the tree. To find the diameter of a tree, we can do a Depth First Search from an arbitrary connection point, find the farthest connection point, and do another Depth First Search from that point. The length of the longest path from that connection point would be the diameter of the tree.

Time Complexity: ~\mathcal{O}(N)~

For subtask 2, we can store all values that satisfy the constraint ~x^2+x=y_i~ for some integer ~x~. We can loop from ~x=1~ to ~x=\sqrt{10^6}~ and store all ~x^2+x~ in a Balanced Binary Search Tree. Then, we can check whether a connection point is valid in logarithmic time by querying if that value exists in the BBST. To find the longest good segment of the tree, we can find the diameter of each "component" of the tree, where each "component" only consists of valid nodes.

Time Complexity: ~\mathcal{O}(N + \sqrt{\max(y_i)}\log\sqrt{\max(y_i)})~

For subtask 3, we must find a quick way of checking whether the constraint is satisfied, as storing ~\sqrt{10^{16}}~ values in a BBST is too slow. There are multiple ways of checking, and we will describe a few here.

One way is by performing a binary search on the value of ~x~ for a certain ~y_i~.

Another way is by checking whether ~x^2+x=y_i~ for ~x=\lfloor{\sqrt{y_i}}\rfloor~.

The third way is by using the quadratic formula:

$$\displaystyle \begin{align*}x^2+x &= y_i \\ x^2+x-y_i &= 0\end{align*}$$

We can plug the values ~a=1, b=1, c=-y_i~ into the quadratic formula:

$$\displaystyle \frac{-b \pm \sqrt{b^2-4ac}}{2a} = \frac{-1 \pm \sqrt{1+4y_i}} 2$$

If the equation evaluates to an integer, then the constraint is satisfied, otherwise, it is not.

Time Complexity: ~\mathcal{O}(N)~ or ~\mathcal{O}(N\log\sqrt{\max(y_i)})~