There are two aspects of this problem: determining an allocation of students to study groups, and determining the optimal adjustment to the plan.

Determining an allocation of students to study groups

The first thing to do is to rephrase this problem under a graph theory framework. We have a functional graph with edges from ~X~ to ~P_X~. Define ~f(X, Y)~ as the node reached after traversing ~Y~ edges, starting from node ~X~. We require that ~f(X, Z) = f(Y, Z)~ implies that ~X~ and ~Y~ are in the same study group.

Recall that every component in a functional graph is a directed cycle, where each node on the cycle is the root of a directed tree. It follows that every student will spend the first few days moving along one of these trees, and once the have reached the root of the tree, they will move around a cycle indefinitely.

Since students on the same node must belong to the same study group, the number of study groups is bounded by the sum of cycle sizes. As it turns out, this bound is always achievable; assign student ~X~ to the ~f(X, N)~-th study group. (We can use some other big number instead of ~N~, such as ~2^{20}~.) We can obtain these values using by computing, for each node, the distance to a cycle and the length of that cycle. We could alternatively implement binary lifting.

Determining the optimal adjustment to the plan

We are allowed to change the successor of a single node, and want to maximise the sum of cycle sizes. How can we do this?

After some thought, the biggest increase to the sum of cycle sizes is the maximum depth of any tree. Consider the tree with the maximum depth, and assume it has more than one node (otherwise the sum of cycle sizes cannot be increased, and it suffices to choose ~A = 1, B = P_1~). Let ~X~ be the root of this tree, ~Y~ be the deepest leaf, ~Z~ be the child of ~X~ which contains ~Y~ in its subtree, and ~W~ be the node which is part of the same cycle as ~X~ with ~P_W = X~. Then, there are two constructions:

• Choose ~A = Z~ and ~B = Y~, creating a new component.
• Choose ~A = W~ and ~B = Y~, increasing the size of the cycle in the current component.

Time complexity: ~\mathcal{O}(N)~ or ~\mathcal{O}(N\log{N})~