SUMSUBARR - 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:

TBD

PREREQUISITES:

Prefix sums, stacks, binary search

PROBLEM:

For an array A with non-zero elements, define:

  • f_0(A) = \text{sum}(A)
  • For 1 \le i \le N,
    • If A_i \gt 0, f_i(A) = \sum_{j \lt i} |A_j| - A_i + \sum_{j \gt i} A_j
    • If A_i \lt 0, f_i(A) = -10^{18}
  • f(A) = \max_i (f_i(A))

Given an array A, compute the sum of f(A[L, R]) across all 1 \le L \le R \le N

EXPLANATION:

To solve this problem, it’s helpful to start by first rewriting the expression for f(A) to use a fewer number of variables.
Since we’re dealing with range sums on an array, this means prefix sums are naturally helpful.

Let’s define:

  • P_i = A_1 + \ldots + A_i to be the i-th prefix sum of A.
  • Q_i = |A_1| + \ldots + |A_i| to be the i-th prefix sum of absolute values of A.

With this, we have:

  • f_0(A) = P_N
  • If A_i \gt 0 then f_i(A) = Q_{i-1} - A_i + P_N - P_i

So, if we then define B_i = Q_{i-1} - A_i - P_i if A_i \gt 0, and B_i = -10^{18} otherwise, we have

f(A) = P_N + \max(0, \max_i B_i)

It’s then not too hard to turn this into an expression for computing answers for subarrays as well.
To find the answer for A[L, R], we have:

  • A base value of P_R - P_{L-1}
  • For each index i \in [L, R] with A_i \gt 0, a value of Q_{i-1} - Q_{L-1} - A_i + P_R - P_i.
    This is simply B_i - Q_{L-1} + P_R.
  • So, we have f(A[L, R]) = P_R + \max(-P_{L-1}, (\max_i B_i) - Q_{L-1})

In particular, note that the quantities P_i, Q_i, B_i are the ones computed for the entire array; we didn’t recompute anything specifically for [L, R].

Our goal is hence to find the sum of P_R + \max(-P_{L-1}, (\max_i B_i) - Q_{L-1}) across all 1 \le L \le R \le N.

Summing up P_R across all such pairs is trivial, so we focus on the second term only.

We have the maximum of two elements, and one of those two elements is itself the maximum of several values.
To deal with this, it’s helpful to try and fix quantities.

So, let’s first fix the index k such that B_k = \max_i B_i for the subarray we’re considering.
We’re now interested in subarrays [L, R] that:

  1. Contain index k, and
  2. Don’t contain another index i such that B_i \gt B_k.

To avoid overcounting, we allow elements equal to B_k to the right but not to the left.
Thus, if x_k \lt k is the largest index such that B_{x_k} \ge B_k, and y_k \gt k is the smallest index such that B_{y_k} \gt B_k, we’re currently considering exactly those subarrays for which x_k \lt L \le k \le R \lt y_k.

Our goal is to compute the sum of \max(-P_{L-1}, B_k - Q_{L-1}) across all of these subarrays.
Note that this depends only on the value of L.
Also, we have -P_{L-1} \ge B_k - Q_{L-1} \iff Q_{L-1} - P_{L-1} \ge B_k.

So,

  • For each index L \in [x+1, k] satisfying Q_{L-1} - P_{L-1} \ge B_k, we want to add -P_{L-1}
  • For every other index L, we want to add B_k - Q_{L-1}

The key observation here is that the values of (Q_{L-1} - P_{L-1}) are non-decreasing with L.
This is because positive elements of the array don’t affect this difference at all; while negative elements add twice their absolute value to it.

Thus, the indices L satisfying Q_{L-1} - P_{L-1} \ge B_k will form a suffix of the range [x+1, k].
Finding the appropriate suffix is easily done with binary search, after which:

  • Computing the requisite sum of P_{L-1} is easy - it’s a range sum over the P array so simple prefix sums will do.
  • Computing the requisite sum of (B_k - Q_{L-1}) is similarly easy: B_k is a constant so really this is just a range sum over the Q array, once again prefix sums do the job.

Thus, with a fixed index k, as long as we know the border indices x_k and y_k, we can compute its contribution to the sum in \mathcal{O}(\log N) time, needing only one binary search and a couple of prefix sum lookups.

Computing the indices x_k and y_k can be done simultaneously for all k in linear time using a stack.

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()))

    p, q = [0], [0]
    ppre, qpre = [0], [0]
    b = [0]
    for x in a:
        p.append(p[-1] + x)
        ppre.append(ppre[-1] + p[-1])
        
        if x > 0: b.append(q[-1] - x - p[-1])
        else: b.append(-10**18)

        q.append(q[-1] + abs(x))
        qpre.append(qpre[-1] + q[-1])
    
    lt, rt = [0]*(n+2), [n+1]*(n+2)
    stk = [0]
    for i in range(1, n+1):
        while stk[-1] > 0:
            if b[stk[-1]] < b[i]: stk.pop()
            else: break
        lt[i] = stk[-1]
        stk.append(i)
    stk = [n+1]
    for i in reversed(range(1, n+1)):
        while stk[-1] <= n:
            if b[stk[-1]] <= b[i]: stk.pop()
            else: break
        rt[i] = stk[-1]
        stk.append(i)
    
    ans = 0
    for k in range(1, n+1):
        ans += p[k]*k

        x, y = lt[k], rt[k]
        lo, hi = x, k
        while lo < hi:
            mid = (lo + hi)//2
            if q[mid] - p[mid] >= b[k]: hi = mid
            else: lo = mid + 1
        
        val = b[k] * (lo - x)
        val -= ppre[k-1] - ppre[max(0, lo-1)]
        val -= qpre[max(0, lo-1)] - qpre[max(0, x-1)]
        ans += val * (y - k)
    print(ans)