PROBLEM LINK:
Author: Vivek Hamirwasia
Tester: Shiplu Hawlader
Editorialist: Praveen Dhinwa
DIFFICULTY:
Medium
PREREQUISITES:
number theory, binary search
PROBLEM:
You are given an array A. You are given Q queries, In each query, you are given a number x and asked to find f(A, x) where
f(A, x) = root(x, 1) * A[1] + root(x, 2) * A[2] + … + root(x, N) * A[N].
root(N, x) is defined as x th root of N. In more formal terms, it is the largest y such that y x <= N.
QUICK EXPLANATION
-
Most important observation is that the root(x, k) will decrease very fast and will become 1 soon.
If at some iteration, value of root has become 1, then we don’t need to iterate over the entire remaining part of the array,
we simply need to find sum of the remaining array.
For finding sum of a specific part of array, we can use calculate prefix sums and use them to compute the sum of any part.
eg. Let prefixSum[i] denotes the sum of array from 1 to i. Then sum of part of array from L, R can be easily found by prefixSum[R] - prefixSum[L - 1].We can simply note that in our case, we needs sum of the suffixes of the array (ie R = N). So we can simply compute sum of suffixes
and directly use them to calculate our answer. -
Root(x, 60) = 1 for all x ≤ 2 60 .
-
Hence for each q from 1 to 60, we need to find p such that p q ≤ x.
-
We can do a binary search to find value of p.
-
In the binary search, we need to check the conditions of form p q ≤ x. We can not do that by simple looping in O(q)
time as it will get TLEd.
So we need to check the above given condition faster. For that we can make use of following facts.-
if q = 1, then p = x.
-
if q = 2, then p = sqrt(x).
-
if q = 3. then p <= 10 6 . So we will already store all possible powers of numbers <= 10 6 .
-
-
After doing these optimizations, we will make sure that p is surely less than 10 6 . Now we can do a binary search of p over
values from 1 to 10 6 .
Note that we need O(1) computation for computing the powers of numbers below 10 6 as they are already stored in the pre-computation
step mentioned above.
EXPLANATION
We will show that root(x, k) decreases very fast and becomes 1 in at most 60 steps. So for x ≤ 10 16 , root(x, k) will reduce to 1
in at most 60 steps.
Note that root(x, k) means to find largest p such that p k ≤ x.
For x = 2 60 and k = 60. p will be 2.
For x = 2 60 and k = 61, p will be 1.
If k becomes largest than log(x) then p will become 1. Hence we have proved this fact.
So by using this fact, we simply need to find first 60 terms only.
Now how to find f(A, x) ?
Note f(A, x) = root(x, 1) * A[1] + root(x, 2) A[2] + … + root(x, N) * A[N].
So f(A, x) = root(x, 1) * A[1] + root(x, 2) A[2] + … + root(x, 60) A[60] + 1 * A[61]+ 1 * A[N].
Hence it becomes f(A, x) = root(x, 1) * A[1] + root(x, 2) A[2] + … + root(x, 60) A[60] + Sum of the array from 61 to N.
Finding the sum of the array from 61 to N, you can store the cumulative sum up to i th index and make use of that.
You can also store the suffix sum ie suffixSum[i] will denote sum of elements from i to N. Hence sum of array from 61 to N will be suffixSum[61].
So will iterate over q from 1 to 60. We need to find largest p such that p q <= x.
Note that we can just iterate over values of p from 1 to x and check whether p q <= x. Complexity of this solution is too high, because it
will need O(x) step, Our x is at most 10 18 , doing 10 18 computation won’t fit in the time.
So we need to somehow optimize this.
Note that p^q is an non-decreasing function over q.
So we can do binary search on the function.
Can we somehow get a good estimate of values of p, Yes, we can use pow function in C++/Java and get pow(x, 1.0 / q), it will give us a good estimate of
possible values of p, as the computation in double is not reliable, we will take some margin of error (let us say 100).
So our range for binary search will be pow(x, 1.0 / q) - 100 and pow(x, 1 .0 / q) + 100.
Now in the binary search routine, we need to check whether given p, q, x whether p q <= x or not?
We can do this by simply powering directly in O(q) time. But it will time out in our case.
We can compute this by powering by squaring in O(log(q)) time, Unfortunately, this will time out in our case.
We need to do something fast.
We make following simple but important observations.
- if q = 1, then p = x.
- if q = 2, then p = sqrt(x).
- if q > 2. then p <= 10 6 . So we will already store all possible powers of numbers <= 10 6 .
So in the case of q > 2, we will store all the powers of the numbers p, Note that these powers are not much, they can be at
most (10 6 log(10 6 ))
which will easily fit in the memory (Proof of this is left for readers as homework).
So we will first check whether q <= 2 or not. If yes, then we can directly compute the answer using the previous observations.
Otherwise we will do a binary search of p over values from 1 to 10 6 .
Note that we need O(1) computation for computing the powers of numbers below 10 6 as they are already stored in the pre-computation
step mentioned above.
Note the following important point.
If you have to check whether a * b <= x or not, where a, b and x all fit in long long.
Then we can not directly check it by doing if (a * b <= x), because a * b will not fit in long long and will overflow,
which will make our calculations erroneous.
Instead we can check the above condition in the slightly modified way, ie. if (a <= x / b). In this way both a and x / b will fit in long long and
wont overflow. You should pick up this trick and it will be really helpful in programming contests.
Complexity:
O(log(10 6 ) * 60) per query.