Problem Link:
Author Roman Furko
Tester Balajiganpathi
Editorialist Utkarsh Lath
Difficulty:
Medium
Pre-requisites:
None
Problem:
Given an integer q, a number is said to be magic number if sum of square of its digits is at most q. You are given an array A[1], A[2], … A[N], you have to build the array B such that B[i] is A[i]th magic number. Then you have to answer some M queries, which ask for sum of all numbers in the array B from given indices l to r.
Explanation:
The problem essentially asks us to find the ith magic number for any 1 ≤ i ≤ 300000. Once this can be done, we also need to sum of arbitrary sub-arrays, but we will come to that part later.
The simplest way to find the kth magic number is to check for all integers from 1 onwards if they are magic or not, and report the kth integer which is magic. Given an integer i, we can easily test if it is a magic integer by finding sum of square of its digits.
Is this sufficient to solve the problem at hand ? Our solution may end up taking too much time to test all the numbers. We need to find out how many integers our code will test for magic-ness before it finds 300000th magic number. For small values of q, this number could be huge.
For instance, for q = 1, only numbers of the form 10i are magic numbers, so the 300000th number is as huge as 10299999. Similarly, for q=2, the numbers of the form 10i + 10j(and 10i) are magic, so the 300000th number is approximately 10774. It is easy to see that as we increase q, more and more integers become magic integers. In fact, for q=200, the 300000th magic integer is approximately 360000, therefore subtask 1 can be solved using this idea, assuming we will be able to report sum of arbitrary sub-arrays efficiently.
It is clear then, that to solve subtask 2, we will need to generate magic numbers efficiently. To do this, we should first try to analyse what went wrong with our initial method for small values of q.
Lets say q=5, and we are trying to generate magic numbers in ascending order. We would be testing 1, 2, 3 … etc for magic-ness. It is clear that when we reach 3, the sum of square of digits becomes 9, so it is not a magic number. However, our trivial algorithm will continue with checking 4,5 etc for magic-ness. We can see that this is clearly a waste of time because if 3 is not a magic number than any number from 3-9 is also not a magic number, and we could straight-away start our search from 10. In general we can say the following
Given a number n of length l which should be tested for magicness, if the sum of its first(most significant) t digits is more than q, then
- n is not a magic integer.
- Any other integer of length l whose first t digits are same as n cannot be a magic number.
Therefore, we could skip all such numbers, and continue testing from the first number which does not satisfy the above property.
The above idea can be implemented very crisply using recursion as shown below. I have used the fact that we can generate all numbers of length l by first fixing the first digit to something between 1 and 9, then fixing the second digit to something between 0 and 9, and so on.
int magic-numbers[300001], count = 1;
// generate _all_ magic numbers of length **l**, whose first **prefix-len** digits have been fixed, and the sum of squares of these digits is **sq-sum-so-far**. **prefix-value** is the integer represented by these **prefix-len** digits.
void gen-all-of-length-l(int l, int sq-sum-so-far, int prefix-len, int prefix-value)
if q == sq-sum-so-far || prefix-len == l
// all remaining digits must be zero to generate a magic number
magic-numbers[count++] = prefix-value * pow(10, prefix-len - l)
return
for (dig = 0; dig * dig + sq-sum-so-far ≤ q; dig++)
// set **prefix-len + 1<sup>th</sup>** digit as **dig**.
gen-all-of-length-l(l, dig * dig + sq-sum-so-far, prefix-len+1, (prefix-value * 10 + dig)%1000000007 )
for(len = 1; count <= 300000; len++)
// generate all magic numbers of length **len**
for dig = 1 to 9
// the first digit is set too dig
if dig * dig ≤ q
gen-all-of-length-l(len, dig * dig, 1, dig)
This solution can find first k magic numbers in time O(k).
Proof:
It is easy to see that if the largest magic number we generate is of length len, then the total number of recursive calls is at most O(len * k), because for every magic number generated, at most len recursive calls are made.
During recursion, when a function is called, it either
- Generates a single magic integer and exits, or,
- Makes two or more recursive calls[dig loops at least upto 1].
Lets call recursive calls of type 1 good calls, and calls of type 2 bad calls. We know that number of good calls is k, so we need to restrict the umber of bad calls. If bj is the number of bad calls for which prefix-len = j, and gj is the number of good calls for which prefix-len = j then:
bj ≤ (gj+1 + bj+1)/2
because every bad call with prefix-len=j makes at least two recursive calls with prefix-len = j+1. Summing the above over all 1 ≤ j ≤ len, we get
B ≤ (G - g1 + B - b1)/2
⇒ B ≤ G - g1 - b1 ≤ k
Where G = sumj≥1 gj is the total number of good calls and B = sumj≥1 bj is the total number of bad calls.
Reporting sums of subarray queries
We need to find sum of elements of array B from index l to r. If we do this naively in O(r-l) = O(N) time, we could end up with time complexity of O(N * M). For the problem constraints, this would take approx. 3 * 1010, which is too much. Therefore, doing this naively will would destroy all our hard work so far. There must be an easy way to report the subarray sums in small time, otherwise we are stuck here. And indeed, there is a way. If we store another C defined as
C[i] = Σ j=1i B[j]
Then we could report sum of the subarray from l to r in O(1) time as C[r] - C[l-1](% 1000000007).