Divide And Conquer, Number Theory
Define a beautiful number as one whose sum of divisors squared is odd. You are given a tree, in which each node has a number. A path from node X to node Y is beautiful iff the number of beautiful nodes is atleast the number of ugly (non-beautiful) nodes along the path. The length of a path is the number of nodes on the path. Find the total length of all beautiful paths over all pairs (X, Y) of nodes. Do not consider (X, Y) distinct from (Y, X) in this count.
Determining whether a number is beautiful or not
Let us first find out how to determine whether a number is beautiful or not. So, given the number n, what is the value of the sum of its divisors squared? Generally, most functions that are sums of functions of the divisors are arithmetic functions (i.e. f(mn) = f(m)f(n) when m and n are coprime). You could make this guess for this function too, and then prove it.
So, if n = pk, then its divisors are 1, p, p2, . . ., pk. The sum of its divisors squared is now 1 + p2 + p4 + … + p2k. If p == 2, then this the sum of these numbers is definitely odd. If p was an odd prime, then the sum of these numbers is odd iff k is even.
Now, when n = ∏(piai). The sum of squares of piai is ∑j=0ai (pi2j). Product of all these when expanded gives you exactly one square term for each divisor. Hence, f(n) = prod(f(piai)).
Now, f(n) is odd iff f(piai) is odd. This means that for every odd divisor of n, its exponent must be even for n to be beautiful.
By prime factorizing numbers upto 1.5 * 10^6, one can check in O(1) time whether a given number is beautiful or not.
Finding the total length of all beautiful paths
Now, we are left with finding the total length of all beautiful paths.
Since this is a tree, let us pick a “root” vertex R and check what is the total length of beautiful paths passing through R. We then apply Divide and Conquer: by removing R and considering the subproblems of each remaining connected component separately, and add all the answers up.
So, we have now rooted the tree at R, and each node is either beautiful or not. Let R’s children be x1, x2, …, xk.
Now, suppose you are at a node y in the tree. Ask yourself what is the sum of lengths of beautiful paths which start at y, and pass through R?
It is easy to find the beauty of the path from R to each node in the tree (inclusive of both). You are currently at y, which has a path-beauty of b(y). Let us say the child of R which is y’s ancestor is xi. You need to ask, what all nodes are there in the subtrees rooted at x1, x2, …, xi-1, whose path-beauty value is >= -b(y)+b®. This is because, any path from y to such a node would go through R, and would have a total beauty of b(y) + b(other-node) - b®.
What you really need is sum of the lengths, and the number of such nodes, not quite “which nodes”. Therefore, let us keep an abstract array len in which we wish len[b] = length of paths “seen so far” whose beauty = b. By “seen so far”, I mean that if you are node y which is a descendant of xi, then the values stored only correspond to nodes whose ancestors are x1, x2, …, xi-1. Now, you are at node y whose beauty is b(y), and you need to ask, what is sum len[j] for j >= -b(y) + b®. We similarly have an abstract array cnt where cnt[b] = the number of nodes whose beauty is b.
The answer to the question “what is the sum of lengths of beautiful paths which start at y, and pass through R?” is “sum len[j] for j >= -b(y)+b®” + (length of path from y to R, exclusive of R) * “sum cnt[j] for j >= -b(y)+b®”. What I have done is just counted the length of each path by breaking it up as a path from R to the other node + a path from y to R.
This can easily be accomplished by constructing a BIT over the abstract arrays len and cnt. Thus, so far we get the following algorithm. Assume that when you are at a node which is a child of x1, that R’s beauty value is already seen.
ans = 0 Root the tree at R // at this step, compute depths and beauties of all nodes BITLen.increment(-b(R), 1) BITCnt.increment(-b(R), 1) if(b(R) == 1) ans++ // the path (R, R) is beautiful, of length 1 for i from 1 to k ans += do(k) return ans do(xi) ret = 0 For each node y in subtree at xi, ret += BITLen.query(-b(y) + b(R)) + (len(y)-1)*(BITCnt.query(-b(y) + b(R)) // Now, insert all the y's values into BITLen and BITCnt. For each node y in subtree at xi, BITLen.increment(b(y), len(y)) BITCnt.increment(b(y), 1) return ret
We are now nearly done. After we have solved the problem for R, we now remove R and solve the problem for each of the remaining connected components.
If we were to analyze the time complexity of a single one of these runs, we would get O(n log n). If we were to arbitrarily choose R however, then for cases like the tree being a path, if we may end up taking time nlogn + (n-1)log(n-1) + (n-2)log(n-2) + … Which is atleast as bad as O(n^2).
The problem is at each step, the subtree size is very large. If we run a dfs to count the size of subtrees, then we can choose the vertex R intelligently. This approach has been used by the setter.
The overall time complexity is now O(N log^2 N)
Can be found here
Can be found here