### PROBLEM LINK:

### DIFFICULTY:

Medium-hard

### PRE-REQ:

DFS, trees

### PROBLEM:

You are given a rooted tree of N nodes, where root node and some nodes are special. We define \textrm{dist}(u) as distance between node u and its first special ancestor node. We define P as sum of \textrm{dist}(u) for all special nodes u. We need to change parent atmost one node such that tree topology is preserved and P is maximised.

### EXPLANATION:

================

We’ll iterate over all nodes u and assume that we are going to change parent of node u. Obviously new parent cannot lie in subtree of node u(to preserve tree topology); and also a good observation is that the new parent won’t be a special node since that cannot bring a positive change in P. So, we’ll choose best option outside its subtree and calculate the change in P by this change in parent of u.

To calculate these things quickly, we need to maintain two things about each subtree. Now, we define \textrm{open}(u) as number of special nodes in subtree of u, for which first special ancestor would be parent of u. Also, we define \textrm{sum}(u) as the value of P for subtree of u(*i.e.* value of P if tree was only subtree of u and node u is special).

Now, let’s assume we have these values somehow. How do we calculate change in value of P by changing parent of node u to v. The loss is \textrm{open}(u)*\textrm{dist}(u) + \textrm{sum}(u), since the path length from node u to its first special ancestor is being contributed by each of the \textrm{open}(u) special nodes. Also, the gain by joining u with v is \textrm{open}(u)*(\textrm{dist}(v)+1) + \textrm{sum}(u), since the path length from node v to its first special ancestor will be contributed by each of the \textrm{open}(u) special nodes. Note that node v should lie outside the subtree of node u, so we can maintain a range maximum structure which can handle range maximum queries in O(\textrm{log}N). So, we choose a node v such that \textrm{dist}(v) is maximised and node v is not special.

You should now observe that term \textrm{sum}(u) is constant in gain and loss and hence needs not be accounted for. To calculate \textrm{open}(u) and \textrm{dist}(u), we just have to use a simple DFS.

So, total complexity is O(N*\textrm{log}N).