INV1 - Editorial

editorial
, , , , ,
1

PROBLEM LINK:

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

Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Segment trees, lazy propagation, stacks, binary search

PROBLEM:

For an array A, define f(A) to be the maximum value of A_i + A_j across all i \lt j such that A_i \gt A_j.
If no such pair exists, f(A) = 0 instead.

Given an array A, compute the sum of f(C) across all contiguous subarrays C of A.

EXPLANATION:

To solve this problem, it is clear that we need to find some property that the answer will satisfy and then figure out how to apply that property to many subarrays simultaneously.

So, let’s start by understanding when exactly some pair can be optimal (or more specifically, when it won’t be optimal.)

Suppose we choose i \lt j such that A_i \gt A_j.
Let’s look at all the indices between them, i.e. all k such that i \lt k \lt j.

  • First, if such a k satisfies A_k \ge A_i, then we could’ve instead chosen (k, j) which is still an inversion but with not-smaller sum.
    So, in an optimal solution we can assume no such k exists.
  • Now that A_k \lt A_i for all inbetween k, observe that if we have A_i \gt A_k \ge A_j for any such k, we could’ve chosen (k, i) instead; again obtaining an inversion with not-smaller sum.
    So, once again we can presume no such k exists in an optimal solution.

Thus, there surely exists an optimal pair (i, j) such that all elements between the indices is strictly smaller than A_j.

This tells us that A_i is the nearest greater element to the left of A_j.

For each index i, let’s define l_i \lt i to be the largest index such that A_{l_i} \gt A_i.
So, l denotes the index of the nearest greater element to the left, for each element of A.
This can be found for all indices i in linear time using a stack.

From our initial understanding of the problem, the optimal solution for the whole array will surely come from one of the pairs (l_i, i).
Observe that this in fact doesn’t apply to just the whole array, but every subarray as well!
This is because the “nearest greater element” remains the same when we move to a subarray - with the only caveat being that said element might not exist in the subarray depending on the left border.

In any case, we thus have at most \mathcal{O}(N) pairs of elements that can possibly be the answer, for any subarray.

We can use this property to solve for summing across subarrays.

Let’s sweep across right endpoints of the subarray, R = 1, 2, 3, \ldots, N in order.

Let b be an array such that b_L denotes the answer for the subarray [L, R] where R is the current right endpoint.
When we move the right endpoint R \to R+1, we’ll try to keep the array b updated appropriately; if we’re able to do this we can then add b_1 + b_2 + \ldots + b_{R+1} to the answer to account for all subarrays ending at R+1.

So, how do we perform this update?
Well, observe that when moving to R+1, there’s only one new pair that can possibly affect the answer of any subarray - that being (l_{R+1}, R+1), i.e. pairing R+1 with the nearest greater element to its left.

Let S = A_{R+1} + A_{l_{R+1}} be the sum of this single new pair.
Observe that the value of b_L doesn’t change for all indices \gt l_{R+1}, since they don’t include both elements of the pair anyway.
So, we only need to work with L \le l_{R+1}.

For each such index, we really want to replace b_L with \max(b_L, S) to account for the new pair.
While that is a hard problem in general (though still doable), in this case we have some structure that we can utilize.

Observe that in our case, the array b is monotonically non-increasing, i.e. we have b_L \ge b_{L+1} for all L.
This is obvious: subarrays starting earlier have more options to work with.

What this means for us, is that the indices L for which b_L \ge S will form some prefix.
So, suppose we find the largest index x such that b_x \ge S.
Then,

  • For all L \le x, we want to leave b_L unchanged since S doesn’t improve their answer anyway.
  • For x \lt L \le l_{R+1}, we want to set the value of b_L to S since all of these can utilize the new pair but have current answers smaller than S.
  • For l_{R+1} \lt L we want to leave b_L unchanged since they can’t utilize the new pair.

So really, all our update boils down to, is setting the value S on some range of the array b.
This can be done quickly using a segment tree with lazy propagation!
This segment tree also allows us to quickly find the sum of a range of elements after an update, which is exactly what we need to update the answer.

Note that finding the breakpoint x is not too hard either once we have this segment tree - for example we can just binary search for it and use the segment tree to query for a point value to find x in \mathcal{O}(\log^2 N) time.
It’s possible to improve this to \mathcal{O}(\log N) time using “segment tree descent” which bakes the binary search into the segtree query, if you wish to do so (/your template supports this.)

In any case, since we’re able to quickly perform updates when moving from R \to R+1, simply do this for each R successively to obtain a solution that’s either \mathcal{O}(N\log^2 N) or \mathcal{O}(N\log N) depending on implementation; either of which will work quickly enough for the constraints.

TIME COMPLEXITY:

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

CODE:

Editorialist's code (C++) 
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());

struct Node {
	using T = ll;
	T unit = 0;
	T f(T a, T b) { return a+b; }
 
	Node *l = 0, *r = 0;
	int lo, hi;
	T mset = 0;
	T val = unit;
	Node(int _lo,int _hi):lo(_lo),hi(_hi){}
	T query(int L, int R) {
		if (R <= lo || hi <= L) return unit;
		if (L <= lo && hi <= R) return val;
		push();
		return f(l->query(L, R), r->query(L, R));
	}
	void set(int L, int R, T x) {
		if (R <= lo || hi <= L) return;
		if (L <= lo && hi <= R) {
			mset = x;
			val = (hi-lo)*x;
		}
		else {
			push(), l->set(L, R, x), r->set(L, R, x);
			val = f(l->val, r->val);
		}
	}
	void push() {
		if (!l) {
			int mid = lo + (hi - lo)/2;
			l = new Node(lo, mid); r = new Node(mid, hi);
		}
		if (mset)
			l->set(lo,hi,mset), r->set(lo,hi,mset), mset = 0;
	}
};

int main()
{
    ios::sync_with_stdio(false); cin.tie(0);

    int t; cin >> t;
    while (t--) {
        int n; cin >> n;
        vector a(n, 0);
        for (int &x : a) cin >> x;
        ranges::reverse(a);

        stack<int> st;
        vector nxt(n, n);
        for (int i = n-1; i >= 0; --i) {
            while (!st.empty()) {
                if (a[st.top()] <= a[i]) st.pop();
                else break;
            }
            if (!st.empty()) nxt[i] = st.top();
            st.push(i);
        }

        ll ans = 0;
        Node *seg = new Node(0, n);
        auto upd = [&] (int i, int x) {
            // update answer on i...n with x
            if (seg -> query(i, i+1) >= x) return;

            int lo = i, hi = n-1;
            while (lo < hi) {
                int mid = (lo + hi + 1)/2;
                if (seg -> query(mid, mid+1) < x) lo = mid;
                else hi = mid-1;
            }

            seg -> set(i, lo+1, x);
        };

        for (int i = n-1; i >= 0; --i) {
            if (nxt[i] < n) upd(nxt[i], a[i] + a[nxt[i]]);
            ans += seg -> query(i, n);
        }
        cout << ans << '\n';
    }
}