hello guys , so i was trying to solve a question and came up with this one.
Given a tree with N nodes , you need to find the number of Pth cousins . where a Pth cousin is defined as a pair of nodes x,y and there exists a node z which is Pth ancestor of both x and y . you need to print an array of n integers each denoting the count of ith cousins , where 1<=i<=n .
1<n<=100000
You can easily find Pth ancestor of a single node using binary lifting in O(log(N)) time. Let's for our node U call it L. So now we can find for every node U how many other nodes have L as parent. In fact we can store for every subtree how many nodes there are on distance exactly P. This can be done in O(N) time with precompute.
So now answer will be:
sum(cnt[L] — 1) / 2 for every L (Pth parent of every U)
overall complexity O(N + NlogN).
This solution is good for calculating cousins for one given P . but we actually need solution for every P , 1<=P<=N
Then the best solution I thought of is with complexity O(NsqrtN).
First we represent the tree like an array such that the subtree of node U is represented from ST[U] to EN[U]. This means that subtree of node U contains subtrees of all its children as an subarray of our current subarray.
Lets fix some node U to be the current Pth cousin.
This means we need to find the number of pairs of nodes in the subtree U of such that their distance to U is equal. We compute this answer for every U and then our "ANSWER" to the problem is the sum of those answers.
Lets store in our array for every node, its distance from the root. Lets again use our fixed node U. Distance from any node in U's subtree is equal to dist(root, V) — dist(root, U). But dist(root, U) is a constant so we can just use dist(root, V) which we have already precomputed.
No the problems is like this:
Given an array (our array of the tree) and N queries (ST[U], EN[U]). Find sum of the answers of the queries, where an answer for each query is defined as sum(cnt[distance] * (cnt[distance] — 1) / 2) for every distance. Here cnt[distance] is the number of nodes with distance distance from the root. But in fact the distance is what we will have in our array. This can be done using MO's algorithm in O(NsqrtN) time.
Here is this blog for some more explanation.
BTW, The same idea can be implemented in O(NlogN) with a persistent segment tree :)
Thanks man ! the solution is interesting :)
What's up with purples and over complicated data structures?
Suppose dist[x] is the distance between node x and the root. Then every pair with the same dist is a Pth cousin for some P, so if you want the total count just sum (count[distance] * (count[distance]-1)/2) for all distances.
I leave the case in which you want separate answers for each P as an exercise to the reader :P
I think we are talking about different problems, because in your solution you count each pair of nodes with just one P (equal to the current distance).
So you mean every pair that is a Pth cousin can also be a Qth cousin for all Q>P? If that's the case, you don't need many changes to the algorithm.
For every possible distance, make a compressed tree containing the nodes with that distance (say, k) and the lcas of those nodes (at most k-1). For every node in this compressed tree, you know all leaves of the subtree of that node are Pth cousins for some P (and all higher Ps up to the root). Just take care not to overcount, by not counting pairs of leaves that are children of the same son.
It's possible to implement this in O(n), but the O(n log n) implementation should be simple.
This solution is really smart and simple. Now I feel like I overkilled the problem a lot :D