INTEREST - 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, dynamic programming

PROBLEM:

For an array A, f(A) is defined as follows:

  • Start with \text{score} = 0 and no items with you.
  • For i = 1, 2, 3, \ldots in order,
    • If you have no item, you can take one item with value A_i.
    • Then, if you have one item, you increase \text{score} by its current value, and then discard the item.
    • Finally, if you have an item, its value increases by 1.
  • In the end, you must have no items remaining.
  • f(A) is the maximum possible final value of \text{score}.

Given an array A, compute

\sum_{L=1}^N \sum_{R=L}^N f(A[L, R])

EXPLANATION:

Let’s understand how to compute f(A) for a single array A.

Suppose we pick up the item at index i, and sell it at index j (i \le j).
We then increase the score by A_i + (j - i), since each index where we choose to not sell it increases its value by 1.

Let the items we pick be the ones at indices l_1, l_2, \ldots, l_k, and suppose we sell them at indices r_1, r_2, \ldots, r_k respectively.
Then, there exists an optimal solution that will satisfy the following:

  1. A_{l_i} + (r_i - l_i) \gt 0 for each i.
    That is, each item we pick will be sold for a strictly positive value.
  2. r_i+1 = l_{i+1} for each 1 \le i \lt k.
    That is, after the first item, we’ll always immediately pick the next item after selling the current one.
    This is obvious: if r_i+1 \lt l_{i+1}, we could instead hold item i for longer before selling it and hence obtain a larger profit.
  3. r_k = N, i.e. the last item will be sold on the final day.
    This is for the exact same reason as above.

Now, suppose we only know the value l_1, i.e. the index of the first chosen item.
Observe that for every index i \gt l_1, either the item there will be picked by us (and so contribute A_i to the score), or we will pass over it while already holding an item (and so contribute 1 to the score).
Clearly we can just choose the best option for each index.
Thus, for each i \gt l_1, its contribution to the score will be \max(A_i, 1).

So, if the index l_1 is fixed, the best possible score is
A_{l_1} + \sum_{i \gt l_1} \max(A_i, 1)

f(A) hence simply equals the maximum of this value across all choices of l_1.

However, we can do better than try every choice of l_1.
For this, we use the fact that every item must have a positive score when sold.
The first item we sell must have a positive score when sold.

Let t_i denote the smallest index \ge i such that selling item A_i at index t_i would result in a positive score.
The conditions t_i \ge i and A_i + (t_i - i) \gt 0 together tell us that t_i = \max(i, i - A_i + 1).

Now, observe that the optimal choice of first item is simply to choose whichever item has the smallest value of t_i.
That is, choose whichever item attains a positive score first.
The correctness of this can be shown via an exchange argument fairly easily.

We thus have the following algorithm to compute f(A).

  • Define t_i = \max(i, i - A_i + 1) for each i.
  • Let l_1 denote the index with minimal t_i.
  • Then, f(A) = A_{l_1} + \sum_{i\gt l_1} \max(1, A_i)

This criterion is quite nice, and allows us to easily generalize to computing across all subarrays.

Let’s now move to computing the sum of answers across all subarrays.

Suppose we fix the left endpoint L of the subarray, and want to compute the answers for all R \ge L.
Let l_1 denote the index with minimum value of t_i across all i \ge L.
(Recall that t_i = \max(i, i - A_i + 1)).
Now, observe that:

  • If l_1 \gt L, then solving for [L, R] is equivalent to solving for [l_1, R], since we’ve seen that it’s not optimal to choose elements before l_1 anyway.
  • If l_1 = L, we then have:
    • For R \lt t_L, the answer is 0 since by definition of l_1, no element in this range can attain a positive value till index t_L is reached.
    • For R \ge t_L, the answer is A_L + \sum_{L \lt i \le R} \max(1, A_i).

Now, let’s define dp_L to be the sum of answers of all subarrays starting at L.
We process L in descending order.

Computing the index l_1 which minimizes t_i across all i \ge L is simple (since t_i doesn’t depend on L, and can be precomputed).
If l_1 \gt L, we simply set dp_L = dp_{l_1} and continue.

Otherwise, let m = t_{l_1}, and we have

dp_L = \sum_{R= m}^N \left(A_L + \sum_{L \lt i \le R} \max(1, A_i)\right)

This can be computed quickly by summing up the contribution of each index.
Specifically,

  • A_L is added once for each R, for a total of (N - m + 1) \cdot A_L.
  • For each L \lt i \le m, \max(1, A_i) is added once for each R.
    This reduces to (N-m+1) times some range sum of the \max(1, A_i) values.
  • For each m \lt i \le N, \max(1, A_i) is added (N - i + 1) times.
    This is a suffix sum of the values (N-i+1)\cdot \max(1, A_i).

So, by maintaining suffix sums of \max(1, A_i) and (N-i+1)\cdot\max(1, A_i), the value of dp_L can be computed in constant time.

This allows for all the values of dp to be computed in linear time, and the final answer is their sum.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Editorialist's code (PyPy3) 
for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split())) + [0]
    
    suf1 = [0]*(n+1)
    suf2 = [0]*(n+1)
    for i in reversed(range(n)):
        suf1[i] = suf1[i+1] + max(a[i], 1)
        suf2[i] = suf2[i+1] + max(a[i], 1)*(n-i)
    
    dp = [0]*(n+1)
    loc, rt = n, n
    for i in reversed(range(n)):
        cur_rt = max(i, i - a[i] + 1)
        if cur_rt < rt:
            rt = cur_rt
            loc = i
        
        if loc > i: dp[i] = dp[loc]
        else:
            dp[i] = suf2[rt+1]
            dp[i] += (a[i] + suf1[i+1] - suf1[rt+1]) * (n - rt)
    print(sum(dp) % 998244353)