SUMTWOMAX - Editorial

PROBLEM LINK:

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

Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093

DIFFICULTY:

Easy

PREREQUISITES:

Sets, binary search

PROBLEM:

Given a permutation P, compute the sum of second maximum across all of its subarrays.

EXPLANATION:

Let’s fix an element x, and try to figure out which subarrays have x as their second maximum.
For this to be the case, the subarray must have exactly one element larger than x.

Suppose x appears at index i.
Define l_1 to be the largest index such that l_1 \lt i and P_{l_1} \gt x.
So l_1 is the nearest element to the left of x that’s greater than it.
Similarly define r_1 to be the smallest index such that r_1 \gt i and P_{r_1} \gt x, the nearest greater element to the right of x.

Observe that:

  1. If a subarray contains index i but neither index l_1 nor index r_1, the maximum element of the subarray will be x.
  2. If a subarray contains both index l_1 and index r_1, the second maximum element of the subarray cannot be x (because there are at least two elements larger than it.)

So, x will be the second maximum of a subarray if and only if it includes exactly one of indices l_1 and r_1.


Let’s look at subarrays that contain l_1 but not r_1.

Of course, such subarrays must contain index i in the first place.
Since it cannot contain r_1, the right end of such a subarray must lie in the range [i, r_1-1].

As for the left end, it must certainly be \le l_1.
However, the left end cannot be too small either: in particular, if we choose L to be the left border, then among the elements \{P_L, P_{L+1}, \ldots, P_{i-1}\} there must not be any element larger than x (other then P_{l_1}), or x would cease to be the second maximum.

Specifically, let l_2 be the second-nearest index to the left of x containing a larger element (i.e. l_2 is the largest index \lt l_1 containing a value larger than x.)

Then, the left border of the subarray must be larger than l_2.
So, the left border can only be chosen from the range [l_2+1, l_1].

The choices of left and right border are independent.
So, we have (l_1 - l_2) \cdot (r_1 - i) possible valid subarrays that include l_1 but not r_1, and have x as the second maximum.

Similarly, by defining r_2 \gt r_1 as the second-nearest greater element to the right, the number of valid subarrays containing r_1 but not l_1 can be found out to be (r_2 - r_1) \cdot (i - l_1).

This gives us the total number of valid subarrays with x as the second maximum; so multiply this by x and add it to the answer.

Note that some of l_1, l_2, r_1, r_2 may not exist; so take care of that appropriately.


The only information we really needed was the four indices l_1, l_2, r_1, r_2.

There are several ways to find these quickly.

One way is as follows.
Let S be a set (initially empty).
We’ll process values of x starting from N down to 1. At any instant, S will contain all the indices of elements that are strictly larger than the current x.
So, when looking at the value x,

  • Suppose i is the index of x.
  • The indices l_1, l_2 are just the two largest elements smaller than i present in S.
    These can be found with binary search (if you’re using std::set in C++, the lower_bound function does this.)
  • Similarly, r_1, r_2 are the two smallest elements larger than i in S, which can again be found with binary search.
  • After processing x, insert i into S and continue on.

This gives a solution in \mathcal{O}(N\log N).

There are other methods as well: for example a range maximum query data structure (such as a segment tree/sparse table) along with binary search allows for finding the appropriate indices in \mathcal{O}(\log N) or \mathcal{O}(\log^2 N) time (depending on implementation/choice of data structure) which are also acceptable solutions.

TIME COMPLEXITY:

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

CODE:

Editorialist's code (C++)
// #include <bits/allocator.h>
// #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());

int main()
{
    ios::sync_with_stdio(false); cin.tie(0);
    
    int t; cin >> t;
    while (t--) {
        int n; cin >> n;
        vector p(n, 0);
        for (int &x : p) cin >> x;

        vector pos(n+1, 0);
        for (int i = 0; i < n; ++i) pos[p[i]] = i;

        ll ans = 0;
        set<int> active = {-2, -1, n, n+1};
        for (int x = n; x >= 0; --x) {
            int i = pos[x];

            auto it = active.lower_bound(i);
            int l1 = *prev(it), l2 = *prev(prev(it));
            int r1 = *it, r2 = *next(it);

            if (l1 != -1) ans += 1ll*x*(l1-l2)*(r1-i);
            if (r1 != n) ans += 1ll*x*(r2-r1)*(i-l1);

            active.insert(i);
        }
        cout << ans << '\n';
    }
}
1 Like