This was my approach for this problem which I reckon may not be the most optimal way of solving it. But it worked nevertheless.
If we have the n distances from a node (including itself), we can calculate the kth largest distance. This number - 1 would be our answer for that node.
We fix a root node (say node 1) and calculate the distances of all the nodes from it. Now to find the distances from a child node, we can see that for all the distances in the subtree of the child node, we need to subtract 1 from the distances and for all the distances not in the subtree, we need to add 1 to the distances.
We observe that, if we store in an array the preorder traversal of the tree, each subtree will be a contiguous range in that array. In fact, postorder or inorder traversal would also work, but let’s consider preorder for the sake of this editorial.
Now the task simplifies to this:
Create an array from preorder traversal of the tree and store the start and end array indices for each node of the tree - DFS - O(N)
Calculate the distances of all the nodes from the root node (including itself) - DFS - O(N)
Using DFS, whenever we go from parent node to child node, subtract 1 from the range denoting the subtree of child and add 1 to the remaining values (left range and right range). We will do this using sqrt decomposition on the array. We will discuss more below.
Calculate the Ki th largest value in the complete array. Data structure and algorithm is discussed below.
We will be working on the array that we have created. Let’s take bucket size b. Now partition the array into buckets, b elements in each bucket.
For each bucket we will be storing 2 values. (1) a sorted list of values in that bucket (C++ vector) (2) lazy value = 0, we’ll discuss it’s use soon.
Let’s look at range update. Exactly like sqrt decomposition, for every node completely in the range, we just add or subtract 1 to the lazy value of that bucket. Otherwise we rebuild the particular bucket, for all the values in the bucket we calculate the new value brute force and sort these values. Note: lazy value is not changed in this case.
Time Complexity - O(N/b + blog(b))
For querying the kth largest in the current array, we do this. Let’s consider for the time being that all lazy values are 0. I’ll discuss a binary search approach to get the kth largest element. Let’s say I define a function lessthanx which will tell you the number of items in the whole array less than x. We can calculate this simply by calling lower_bound on all the buckets. This would be O(N/b * log(b))
Now let’s remove the restriction on the lazy value. We can see that the lazy value doesn’t alter the relative positioning of the values inside a bucket, so while querying lessthanx, we can simple call lower_bound on x-lazy to get our answer.
We can apply binary search on x to find the position (or distance) such that it’s the kth largest element in the whole array.
Time Complexity - O(log(N) * N/b * log(b))
We have to do all these in our main DFS,
update: the subtree range -1 and rest of the ranges +1
query on the current array for the required kth largest element
un-update: the subtree range +1 and rest of the ranges -1
Total time complexity - N * (N/b + blog(b) + log(N) * N/b * log(b))
= O(N * (blog(b) + log(N)(N/b)*log(b)))
We can choose the value of b to minimize the time complexity.
Hope this helps! Suggestions are welcome!