PRODEQSUM - Editorial

editorial
, ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Easy

PREREQUISITES:

None

PROBLEM:

Given an array A, count the number of its contiguous subarrays whose product equals its sum.

EXPLANATION:

Often, counting subarrays satisfying a certain property can be done by fixing one end of the subarray and figuring out which choice of other end can work.

So, let’s fix the left end L and try to count all valid right ends R for this left end.

Observe that if we start the right end at R = L and keep moving it right by one step, the sum of the subarray will always strictly increase, while the product will either remain the same (if the new element is a 1), or also increase.

In particular, suppose i and j are two indices such that L \le i \lt j, both A_i and A_j are \gt 1, and all of A_{i+1}, A_{i+2}, \ldots, A_{j-1} are equal to 1.
Then, the product of the subarrays [L, i], [L, i+1], \ldots, [L, j-1] will all be the same, while their sums are increasing by 1 from left to right - which means there can be at most one right endpoint in the range [i, j-1] that has equal product and sum.

With this in mind, a somewhat natural (though slow) algorithm is as follows:

  • Fix a left endpoint L.
  • Let i_1, i_2, \ldots, i_k be the list of indices that are \ge L and contain elements \gt 1.
    Assume i_1 \lt i_2 \lt \ldots \lt i_k.
    To account for every case, we’ll assume that i_1 = L (even if A_L = 1), and i_k = N+1. This is just so we also consider the cases of subarrays with product 1, and those that end after the last non-1 value in the array.
  • Now, for each 1 \le j \lt k, do the following:
    • Let P denote the product of the subarray A[L, i_j].
    • Let S denote the sum of the subarray A[L, i_j].
    • Then, P will be the product for all the subarrays ending at indices i_j, i_j+1, \ldots, i_{j+1}-1.
      The sums of these subarrays will be S, S+1, S+2, \ldots, S + (i_{j+1} - i_j - 1).
    • So, if P lies in the interval formed by the sums, increment the answer by 1; otherwise do nothing.

The above algorithm is pretty obviously correct, since it followed from our observation of the product being constant in some segment meaning that at most one right endpoint in that segment would work.
However, there are currently a couple of issues with this algorithm:

  1. The biggest issue by far is speed: this clearly runs in \mathcal{O}(N^2) in the worst case, which is far too slow for us.
  2. Another issue is with respect to computing products of segments - this needs to be done quickly, and unlike the case of sums, prefix products don’t work as well because they can become extremely large.

Luckily, there’s a single observation that allows us to fix both above issues simultaneously.
The key here is to note that because the product of several elements grows very quickly, any segment with equal sum and product cannot have too many elements \gt 1.

In particular, because A_i \le N, any segment of the array has a sum not exceeding N^2, which for N \le 2\cdot 10^5 gives an upper bound of 4\cdot 10^{10}.
This means we don’t need to consider segments whose products exceed this value at all.
But since 2^{40} \gt 4\cdot 10^{10}, this also means we don’t need to consider any subarrays that contain \ge 40 elements that are \gt 1, since their products will be \gt N^2 for sure!

This allows us to modify our initial algorithm as follows:

  • Fix a left endpoint L.
  • Initialize i = L, P = A_L, S = A_L.
    We’ll always maintain P and S to be the product/sum respectively of the subarray A[L, i].
  • Now, repeat the following process:
    • Let j \gt i be the nearest next position containing an element \gt 1.
      If no such position exists, set j = N+1.
      There are several methods of finding this j quickly - for example you can store a list of positions of elements \gt 1 and binary search on it, or you can just precompute this value for each i by iterating in reverse.
    • We now know that the product is constant (equal to P) on all indices i, i+1, \ldots, j-1, while the sum starts at S and increases to S+(j-i-1).
      So, if S \le P \le S + (j-i-1), add 1 to the answer; otherwise do nothing.
    • Next, we need to update the right endpoint to j.
      This can be done as follows:
      • Multiply P by A_j, since only that changes the product.
      • Add A_j + (j - i - 1) to S, since there are (j-i-1) ones and also A_j later.
      • Then, set i to j.
    • Importantly, you can break out as soon as P \gt N^2, since as we established earlier, no further subarray can be valid since their products will be too large.

The last early-break step is the most important step here - it ensures that for a fixed left index L, we process no more than \log_2(N^2) elements to its right that are \gt 1.
This optimization alone thus brings the complexity down to \mathcal{O}(N\log(N^2)) = \mathcal{O}(N\log N).

Further, we no longer need to worry about products becoming too large - we break as soon as a product becomes \gt N^2, so the maximum possible product we’ll ever need to process is N^2 \cdot N = N^3, which for N \le 2\cdot 10^5 fits well within a 64-bit integer datatype.

TIME COMPLEXITY:

\mathcal{O}(N\log N) per testcase.

CODE:

Editorialist's code (PyPy3) 
for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split())) + [0]
    jump = [n]*(n+1)
    for i in reversed(range(n)):
        jump[i] = jump[i+1]
        if a[i+1] > 1: jump[i] = i+1
    
    ans = 0
    for i in range(n):
        L, prod, sm = i, a[i], a[i]
        while prod <= n*n and L < n:
            req = prod - sm
            if 0 <= req <= jump[L]-L-1: ans += 1
            
            prod *= a[jump[L]]
            sm += jump[L]-L-1+a[jump[L]]
            L = jump[L]
    print(ans)
2

Very nice problem!