Editorial for An Animal Contest 5 P6 - Larry Finally Uses His Magical Powers
Submitting an official solution before solving the problem yourself is a bannable offence.
The optimal path to reach any node never involves going up the tree. Therefore, the answer for any given node is .
When solving this problem, the first thing to note is that it is a shortest path problem. We are working with a graph of nodes such that any two nodes with distance of less than in the tree, have an edge in this graph. All of the edges to a node have cost .
So first, let's think about how we would do this naively. The algorithm that should come to mind is Dijkstra, which can be performed in in a graph of edges and nodes. In this case, can be as large as , so the overall complexity would be which is obviously too slow for .
Now the question is, how can we optimize this approach to be fast enough for ? The key observation that must be made here is that all edges pointing towards a node have the same cost. This means that when we are running Dijkstra with the priority queue, the first time that we are able to reach a particular node, is the only time its answer will be updated, and the only time it will be pushed into the priority queue. This is great news as now the only edges will have to be considered! As soon as a node we are able to reach a node, we update its cost and push it into the priority queue. After that, we can mark it as being visited and ignore it (this is similar to what we do in BFS; nodes are marked visited right away so they are only pushed into the queue once).
But how exactly do we find the paths to all unvisited nodes in a reasonable time complexity? We can do this with centroid decomposition. The key property that we must take advantage of is that the lowest common ancestor (LCA) of two nodes in the centroid tree is on the path between any two nodes in the actual tree. Another important property is that in total, there are at most ancestors of any node in the centroid tree. Combining these two properties, there are at most paths in the centroid tree which we can concatenate to form any of the paths in the actual tree.
With these two properties in mind, we will perform our centroid decomposition. Every single one of the paths we encounter, we will store for use during the Dijkstra. We will also store which centroid each of the paths belongs to, and the distance to the centroid. These paths must be stored in sorted order of distance (the reason will become clear later in the editorial). We can get the paths using DFS and sort them, but this is inefficient as it introduces an extra factor. Instead, we will get all of the paths using BFS which guarantees that they will already be in sorted order.
Finally, we can return to solving the actual problem. We will perform Dijkstra, and every time we are processing a node , we can find all of the paths to valid unvisited nodes using the centroid decomposition. We will climb up the centroid tree and for any at each centroid, we will loop through all of the paths to see which nodes are both unvisited and satisfy the distance requirement. Because the paths at a particular centroid are sorted in order of distance, we can maintain a pointer to the last node we processed at the particular centroid. All of the nodes before the pointer are also guaranteed to have been processed, since the paths are sorted in order of distance. Therefore, if a later path has been processed, then so must the previous one. Note that a node might be processed using another centroid in which case we do not advance the pointer for this centroid. We will simply skip any nodes that are processed using another centroid as we go along. In this way, all of the paths will only be looped over once each, and our total time complexity will be . Thus, the problem has been solved.