PROBLEM LINK:
Author: Vasya Antoniuk
Testers: Sergey Kulik and Kevin Atienza
Translators: Vasya Antoniuk (Russian), Team VNOI (Vietnamese) and Hu Zecong (Mandarin)
Editorialist: Kevin Atienza
DIFFICULTY:
Medium
PREREQUISITES:
Algebra, Series
PROBLEM:
Chef has an array A[1\ldots n]. The following procedure is repeated m times:
- For every two neighboring elements a and b, insert a + b between them.
(This happens simultaneously at each step.)
Given x and y, what is the sum of all values from the x th to the y th index in the final array?
QUICK EXPLANATION:
Let F(x) be the sum of the values from the first to the x th index. Then the answer is F(y) - F(x-1), so we can focus on just F(x).
The size of the final array is (n - 1)2^m + 1. The index of the original A[i] in the final array is (i-1)2^m + 1.
Between any two neighboring elements a and b, there will be 2^m - 1 values inserted between them, and the sum of these values is (a + b)\frac{3^m - 1}{2}. Also, these values can be organized into a binary tree, the root of which is the value a + b. We can compute any prefix sum of the inserted values by appropriately traversing a path in this “tree” in O(m) time, without actually constructing the tree.
To compute F(x):
- For every index i such that i2^m \le x, add the value A[i] + (A[i] + A[i+1])\frac{3^m - 1}{2} to the answer. Stop at the first index i such that i2^m > x.
- Compute the appropriate prefix sum in the “tree” between A[i] and A[i+1].
EXPLANATION:
In this problem, we can’t construct the final array explicitly because it’s very large! To give you an idea, notice that after each step, the number of elements roughly doubles. Thus, the final approximately n2^m elements! To be exact, the final array actually has a size of (n-1)2^m + 1. For n = 10^5 and m = 30, this is approximately 10^{14}. Even constructing all values between two elements of the original array isn’t enough, because there are 2^m-1 elements between them. For m = 30, this is approximately 10^9.
In questions asking you to find a sum of a contiguous subarray, it usually helps to convert it to prefix sums instead. In our case, let F(x) be the sum of the values from the first to the x th index. Then the answer is F(y) - F(x-1). Sure, it costs us two queries instead of one, but the benefit is that the solution is potentially much simpler. So now we can focus on just F(x), i.e. we’ll focus on calculating a prefix sum of the final array.
For every two neighboring elements a and b, notice that the numbers that appear between them aren’t affected by other numbers in the array. So it makes sense to study the numbers that appear between a and b. Let’s denote by V_{a,b}(m) the array of intermediate values between a and b after m steps. For example, F_{1,6}(1) = [7] and F_{1,6}(2) = [8, 7, 13].
Sum of V_{a,b}(m)
First, what is the sum of the values in V_{a,b}(m)? Let’s write this sum as S_{a,b}(m). We have the base case S_{a,b}(0) = 0 because in the beginning, there are no intermediate numbers. Now, let’s try to derive a recurrence.
Suppose V_{a,b}(m) = [v_1, v_2, v_3, \ldots], so that S_{a,b}(m) = v_1 + v_2 + v_3 + \ldots. We want to find a formula for S_{a,b}(m+1). So let’s look at V_{a,b}(m+1). After one step, here’s what the new intermediate numbers look like:
Notice that in this array, every v_i appears exactly three times:
So S_{a,b}(m+1) is three times S_{a,b}(m), except that there is also an a and a b in the sum (at the beginning and end). This gives us the following recurrence
To solve this recurrence, we define an offset function T_{a,b}(m) = S_{a,b}(m) + \frac{a+b}{2} so that S_{a,b}(m) = T_{a,b}(m) - \frac{a+b}{2}. Substituting this above:
For the base case, we have T_{a,b}(0) = S_{a,b}(0) + \frac{a+b}{2} = \frac{a+b}{2}. The recurrence for T_{a,b} is much simpler, because at every step we are only multiplying by 3. It’s easy to see that the closed-form solution for T_{a,b} is:
Thus, we have:
Prefix sum of V_{a,b}(m)
Now let’s answer a harder problem: Given x, find the sum of the first x values in V_{a,b}(m). Again, we can’t construct the array V_{a,b}(m) explicitly. But we can still quickly compute the sum of the first x values once we notice the following structure of V_{a,b}(m):
- The middlemost value is a+b.
- The array of values to the left of a+b is the same as V_{a,a+b}(m-1).
- The array of values to the right of a+b is the same as V_{a+b,b}(m-1).
- The size of V_{a,b}(m) is 2^m - 1.
To see the last one, notice that the size of V_{a,b}(m) only depends on m, not on a or b, so let L_m be this size. Also, notice that |V_{a,b}(m)| = |V_{a,a+b}(m-1)| + |V_{a+b,b}(m-1)| + 1, thus, we have the following recurrence: L_m = 2L_{m-1} + 1. The base case is L_0 = 0. Thus, it’s easy to see that L_m = 2^m - 1 solves this recurrence.
With this observation, we can now devise a divide and conquer algorithm to compute the sum of the first x values of V_{a,b}(m) (assuming 0 \le x \le L_m):
- If m = 0, then the answer is 0.
- If m > 0 and x \le L_{m-1}, then the answer is the sum of the first x values of V_{a,a+b}(m-1).
- If m > 0 and x > L_{m-1}, then the answer is the sum of S_{a,a+b}(m-1), a+b, and the sum of the first x - L_{m-1} - 1 values of V_{a+b,b}(m-1).
Notice that the solution performs one recursive call, but m is smaller by 1, so this halts after m recursions!
def L(m):
return 2^m - 1
def S(a, b, m):
return (a + b) * (3^m - 1)/2
def V_sum(a, b, m, x):
// sum of the first x elements of V(a,b,m)
if m == 0:
return 0
else if x <= L(m-1):
return V_sum(a, a+b, m-1, x)
else:
return S(a, a+b, m-1) + (a+b) + V_sum(a+b, b, m-1, x-L(m-1)-1)
Computing F(x)
We now have all the tools needed to compute F(x).
The final array is actually the following concatenation of arrays:
To simplify this a bit, we do the following things:
- We set A_{n+1} = 0 and add the array V_{A_n,A_{n+1}}(m) at the end. According to problem assumptions, x will be \le the size of the final array, so adding any elements at the end won’t affect the answer.
- Let’s abbreviate [A_i] + V_{A_i,A_{i+1}}(m) as W_i, so our final array expression becomes much simpler:
Notice that:
- The sum of the values in W_i is A_i + S_{A_i, A_{i+1}}(m).
- Each W_i's length is exactly 2^m!
Now, we want to sum the first x values in the array W_1 + W_2+ \ldots + W_n. In this case, the following simple algorithm will do:
def W_sum(i):
return A[i] + S(A[i], A[i+1], m)
def F(x):
ans = 0
i = 1
while x >= 2^m:
x -= 2^m
ans += W_sum(i)
i++
return ans + W_sum(i, x)
// W_sum(i, x) is the sum of the first x elements of W[i]
This simply adds up the contribution of each subarray until we find the first array W_i not completely contained in the prefix, in which case we compute a certain prefix sum of that array (W_sum(i, x)). But remember that W_i is [A_i] + V_{A_i,A_{i+1}}(m), and we already know how to quickly compute prefix sums of anything of the form V_{a,b}(m)! Thus,
def W_sum(i, x):
// sum of the first x elements of W[i]
if x == 0:
return 0
else:
return A[i] + V_sum(A[i], A[i+1], m, x-1)
Combining all the above, our algorithm runs in O(n + m)! (assuming we precompute all powers of 2 and 3 up to the exponent m.) Don’t forget to reduce your intermediate values modulo 10^9 + 7!
Time Complexity:
O(n + m)
