PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Author: Divyansh Rastogi
Tester: Shubham Anand Jain
Editorialist: Nishank Suresh
DIFFICULTY:
Medium
PREREQUISITES:
Stacks, Range minimum/maximum queries, (optionally) Binary search
PROBLEM:
You have an array of length N. You can pick any subarray of size 2 or more of this array, after which Chef removes 2 elements from it. Your score is computed as the sum of the remaining elements.
Maximize your score, given that Chef always chooses elements to minimize your score once a subarray is chosen.
QUICK EXPLANATION:
- Chef always picks the largest and second largest element of the subarray
- Iterate over the position of the second largest element
- Once the second largest is fixed, the largest element is either to its left or its right - try both cases. The positions of these elements can be precomputed with a stack
- Also compute the positions of the second closest larger element to the left/right
- Once all these positions are known, the best answer for a given position can be computed by a couple of range minimum/maximum queries
EXPLANATION:
The very first thing to notice is that once you have chosen a subarray, Chef will always remove the largest two elements from it.
So, our task is to find a subarray such that its sum after removing the largest two elements from it, is maximum.
Of course, this can be done in \mathcal{O}(N^3) by iterating over all subarrays and then finding their maximums, and even optimized to \mathcal{O}(N^2) by smartly keeping track of subarray sums and maximums while iterating. This is too slow for our purposes, and it is not immediately obvious how to optimize it to run any faster.
Let us instead approach the problem differently: suppose we fixed which two elements were being deleted, say A_L and A_R where L\leq R. Then, the maximum sum of a subarray containing these two points can be broken into 3 parts:
- The sum of elements from positions L+1 to R-1
- The maximum sum of a subarray ending at L-1
- The maximum sum of a subarray starting at R+1
However, we are fixing the positions of the maximums - hence, none of the 3 parts can contain any element which is larger than \min(A_L, A_R).
This also leads us to another observation:
Suppose A_L < A_R. Then, R must be the smallest index >L such that A_L < A_R.
Similarly, if A_L > A_R, L must be the largest index <R such that A_L > A_R.
Hence it seems useful to know, for each element, what the next/previous greater element is.
It is well known that computing the next/previous greater element can be done for all indices of an array in \mathcal{O}(N) time using a stack: see this for more information.
Armed with this information, we note something further: upon fixing the second largest element of the array, the largest element is either going to be the next greater or the previous greater.
In other words, we have reduced the number of objects we are iterating on from \mathcal{O}(N^2) to \mathcal{O}(N). If we are able to calculate the maximum subarray sum for each of these (L, R) pairs in something like \mathcal{O}(\log N), that would solve the problem - which is exactly what we will do.
Recall that the subarray was broken into 3 independent parts. Let us try to solve for each of those separately.
The middle part
Upon fixing L and R, every element between them must be taken, no matter what.
So, this simply comes out to be the sum of the subarray between [L+1, R-1] which is easily calculated by, say, prefix sums.
For the remaining two parts, I assume that A_L < A_R - the other case proceeds similarly.
The left part
We want the largest possible subarray sum ending at L-1 such that its elements don’t exceed A_L.
Let L_1 be the index of the previous greater element of index L. For the subarray to not have an element exceeding A_L, it must equivalently not contain index L_1.
So, we are looking for the largest sum of a subarray of the form [i, L - 1] where L_1 < i < L.
In terms of prefix sums, the sum of any such subarray is pref_{L-1} - pref_{i-1} where pref_i = A_1 + A_2 + \dots + A_i and pref_0 = 0.
Note that pref_{L-1} is a constant, so maximizing the sum is equivalent to picking the minimum value of pref_{i-1} for L_1 < i < L.
This is the classical range minimum query problem which can be answered in a host of ways - for our purposes, we will assume \mathcal{O}(\log N) with the help of a segment tree.
The right part
At first glance, this case is symmetric to the one above - if we know the first index larger than R (say R_1) such that A_{R_1} > A_L, the largest subarray sum starting from R+1 and ending before R_1 can be found via a segment tree query again (this time, a range maximum query instead).
What if such an index doesn't exist?
There is no problem in such a case - it just means that every subarray starting from R+1 can be considered, which is equivalent to saying that R_1 = n+1.
However, recall that we have only previously computed the index of the next greater element of R - there is no guarantee that R_1 should also be this index, it might be less.
Example
Consider the array
1 3 2 4
The next larger element of 1 is 3, and the next greater element of 3 is 4. However, the R_1 we are looking for corresponding to 1 is 2, not 4.
So, the problem here is how exactly to find R_1 - if we are able to do that, the problem is functionally solved.
It turns out that there are several ways to do this:
Method 1 (Stacks) (Setter's Solution)
Note that what we want to compute is essentially just the second nearest greater element to the left/right.
To do this, It is possible to extend the algorithm used to compute the nearest greater element.
Here is some pseudocode to compute the next greater element:
stack s
for i from 1 to n:
while s is not empty and a[top(s)] < a[i]:
next_greater[top(s)] = i
pop s
push i onto s
The idea here is that s contains a (monotonic) sequence of indices whose next greater element hasn’t been found yet.
Extending that idea, computing the second next greater element only needs another stack - whenever something is popped from the first stack, it can be pushed into the second, to signify that its next greater element has been found, and whatever pops it from the second stack will be its second next greater element.
However, note that in order to maintain monotonicity of the second stack, we cannot just push values into it as soon as we pop from the first. Instead, at each index, we must store the sequence of elements popped, and then push them in reversed order.
Pseudocode is as follows:
stack s1, s2
for i from 1 to n:
while s2 is not empty and a[top(s2)] < a[i]:
next_greater_2[top(s2)] = i
pop s2
stack rem
while s1 is not empty and a[top(s1]) < a[i]:
next_greater[top(s1)] = i
push top(s1) onto rem
pop s1
while rem is not empty:
push top(rem) onto s2
pop rem
push i onto s1
This still runs in \mathcal{O}(N) time overall because the cost of pushing/popping amortizes over all indices.
Method 2 (Binary search + RMQ) (Tester's and Editorialist's Solution)
it is possible to binary search on the position of R_1, with the help of (yet again) range maximum queries.
Binary search on the range [R+1, N+1], and:
- Suppose the current range is [lo, hi].
- Query for the largest element in [lo, mid], suppose it is x.
- If x > A_L, set hi = mid
- Otherwise, set lo = mid+1
At the end of this, lo is the R_1 we are looking for.
The complexity of this part is \mathcal{O}(\log^2 N) assuming a segment tree is used for the range maximum query, which may or may not pass depending on how well the implementation is handled.
However, it is possible to improve the complexity to \mathcal{O}(\log N) by combining the binary search into the query function of the segment tree. This technique is sometimes called implicit binary search/walking down the segment tree.
Implementation details can be found in the comments of this blog
Implementation Details
The solution outlined above takes some liberty with details in order to not clutter the main idea. In particular, there are some cases where you would like the next/previous greater element to be equal instead of strictly larger, as in the case of the array
1000 -1000000 10 5 5 10 -1000000 1000
Here, the best subarray to choose is 10 5 5 10 giving a final answer of 10, which the above solution doesn’t quite do because it would always pair a 10 with a 1000 which is suboptimal.
The fix
Note that for a given index, we compute two things: the next/previous greater element, and the second next/previous greater element.
By far the simplest way to fix this problem is to compute the first greater element allowing equality, and the second with a strict inequality.
Making this change is just a matter of changing < to <= (or vice versa, depending on implementation) in the popping criteria.
Why does this work?
Consider a case which looks like:
----a------b---x---x--x----c------d-----
where x < a, b, c, d and - represents something smaller than x.
Suppose we fix one of the x between b and c to be the second largest element.
Then, there are several possibilities for the subarray which can be chosen:
- The subarray has to contain at least one of b, c or another x; else the one we fix would not be the second maximum
- The subarray cannot contain both b and c at the same time, and cannot contain a or d no matter what.
- If the subarray contains neither b nor c, it is completely contained between them and will contain some 2 consecutive x's. This case will be covered when fixing the first x as the second largest element, and checking for the largest element on the right (since we allowed equality)
- If the subarray contains b, it must contain the leftmost x. This case will be covered when fixing the first x and checking for the largest element on the left
- If the subarray contains c, it must contain the rightmost x. This case will be covered when fixing the last x and checking for the largest element on the right.
Also, make sure to properly handle cases when the something doesn’t exist (i.e, when there is no larger/second larger element).
TIME COMPLEXITY:
\mathcal{O}(N\log N) or \mathcal{O}(N\log^2 N) per testcase, depending on implementation.
SOLUTIONS:
Setter's Code
# include "bits/stdc++.h"
using namespace std;
template <typename T>
struct SEGTREE {
struct NODE {
T sum;
T pre;
T suf;
};
NODE merge(const NODE &A,const NODE &B) {
NODE C;
C.sum = A.sum + B.sum;
C.pre = max(A.pre, A.sum + B.pre);
C.suf = max(B.suf, B.sum + A.suf);
return C;
}
vector<NODE> t, a;
int n;
void init(vector<T> A) { // initialize on vector
n = A.size();
t.resize(4 * n);
for (auto it : A)
a.push_back(NODE{it, it, it});
build(0, 0, n - 1);
}
void build(int u, int tl, int tr) {
if (tl == tr) {
t[u] = a[tl];
return;
}
int tm = (tl + tr) >> 1;
build(2 * u + 1, tl, tm);
build(2 * u + 2, tm + 1, tr);
t[u] = merge(t[2 * u + 1], t[2 * u + 2]);
}
NODE query(int l, int r) {
if (l > r) return NODE{0, 0, 0};
return _query(0, 0, n - 1, l, r);
}
NODE _query(int u, int tl, int tr, int l, int r) {
if (tl == l && tr == r)
return t[u];
int tm = (tl + tr) >> 1;
if(r <= tm) return _query(2 * u + 1, tl, tm, l, r);
else if(l > tm) return _query(2 * u + 2, tm + 1, tr, l, r);
else return merge(_query(2 * u + 1, tl, tm, l, tm), _query(2 *u + 2, tm + 1, tr, tm + 1, r));
}
};
void solve() {
int n;
cin >> n;
vector<long long> a(n);
for (int i = 0; i < n; ++ i) cin >> a[i];
// INITIATE SEGMENT TREE FOR MAX SUFFIX AND PREFIX
SEGTREE<long long> seg;
seg.init(a);
// CALCULATE NEXT GREATER 1 and 2
vector<int> nge1(n, n), nge2(n, n);
{
vector<pair<long long, int>> st1, st2;
for (int i = 0; i < n; ++ i) {
while (!st2.empty() && st2.back().first < a[i]) {
nge2[st2.back().second] = i;
st2.pop_back();
}
vector<pair<long long, int>> add;
while (!st1.empty() && st1.back().first <= a[i]) {
nge1[st1.back().second] = i;
add.push_back(st1.back());
st1.pop_back();
}
st2.insert(st2.end(), add.rbegin(), add.rend());
st1.push_back({a[i], i});
}
}
// CALCULATE PREVIOUS GREATER 1 and 2
vector<int> pge1(n, -1), pge2(n, -1);
{
vector<pair<long long, int>> st1, st2;
for (int i = n - 1; i >= 0; -- i) {
while (!st2.empty() && st2.back().first < a[i]) {
pge2[st2.back().second] = i;
st2.pop_back();
}
vector<pair<long long, int>> add;
while (!st1.empty() && st1.back().first <= a[i]) {
pge1[st1.back().second] = i;
add.push_back(st1.back());
st1.pop_back();
}
st2.insert(st2.end(), add.rbegin(), add.rend());
st1.push_back({a[i], i});
}
}
long long ans = 0;
// BRUTE FORCE FOR 2nd MAXIMUM
for (int i2 = 0; i2 < n; ++ i2) {
/*
Consider the following, for better visualisation
---- L ------ i1 ------ i2 ------- i3 ------ R ------
i2: is your current 2nd maximum
i1: is first previous greater element of i2
i3: is first next greater element of i2
L: is second previous greater element of i2
R: is second next greater element of i2
*/
int i1 = pge1[i2];
int i3 = nge1[i2];
int L = pge2[i2];
int R = nge2[i2];
// CASE 1
if (i1 != -1) {
long long between = seg.query(i1 + 1, i2 - 1).sum;
int modified_i3 = (i3 < n && a[i3] == a[i2]) ? R : i3;
long long after = max(0LL, seg.query(i2 + 1, modified_i3 - 1).pre);
long long before = max(0LL, seg.query(L + 1, i1 - 1).suf);
ans = max(ans, before + between + after);
}
// CASE 2
if (i3 != n) {
long long between = seg.query(i2 + 1, i3 - 1).sum;
long long after = max(0LL, seg.query(i3 + 1, R - 1).pre);
int modified_i1 = (i1 >= 0 && a[i1] == a[i2]) ? L : i1;
long long before = max(0LL, seg.query(modified_i1 + 1, i2 - 1).suf);
ans = max(ans, before + between + after);
}
}
cout << ans << '\n';
}
int main() {
int tt;
cin >> tt;
while (tt --) {
solve();
}
return 0;
}
Tester's Code
//By TheOneYouWant
#include <bits/stdc++.h>
using namespace std;
#define fastio ios_base::sync_with_stdio(0);cin.tie(0)
#define int long long int
const int LL_MIN = -1e18;
const int LL_MAX = 1e18;
struct SegtreeMax{
vector<int> t;
SegtreeMax(int n) {
t.clear();
for(int i = 0; i < 4*n; i++){
t.push_back(LL_MIN);
}
}
void build(int a[], int v, int tl, int tr){
if(tl == tr){
t[v] = a[tl];
}
else{
int tm = (tl + tr)/2;
build(a, v*2, tl, tm);
build(a, v*2+1, tm+1, tr);
t[v] = max(t[v*2], t[v*2+1]);
}
}
int query(int v, int tl, int tr, int l, int r) {
if(l>r) return LL_MIN;
if(l <= tl && tr <= r){
return t[v];
}
int tm = (tl + tr)/2;
return max(query(v*2, tl, tm, l, min(r, tm)), query(v*2+1, tm+1, tr, max(l, tm+1), r));
}
int query2(int v, int tl, int tr, int l, int r, int val, bool left) {
if(l>r) return -1;
if(l <= tl && tr <= r){
if(left){
if(t[v] <= val) return -1;
}
else{
if(t[v] < val) return -1;
}
if((l==r) && (tl==tr)){
return l;
}
}
int tm = (tl + tr)/2;
if(left){
int a1 = query2(v*2+1, tm+1, tr, max(l, tm+1), r, val, left);
if(a1!=-1) return a1;
return query2(v*2, tl, tm, l, min(r, tm), val, left);
}
else{
int a1 = query2(v*2, tl, tm, l, min(r, tm), val ,left);
if(a1!=-1) return a1;
return query2(v*2+1, tm+1, tr, max(l, tm+1), r, val, left);
}
}
};
struct SegtreeMin{
vector<int> t;
SegtreeMin(int n) {
t.clear();
for(int i = 0; i < 4*n; i++){
t.push_back(LL_MAX);
}
}
void build(int a[], int v, int tl, int tr){
if(tl == tr){
t[v] = a[tl];
}
else{
int tm = (tl + tr)/2;
build(a, v*2, tl, tm);
build(a, v*2+1, tm+1, tr);
t[v] = min(t[v*2], t[v*2+1]);
}
}
int query(int v, int tl, int tr, int l, int r) {
if(l>r) return LL_MAX;
if(l <= tl && tr <= r){
return t[v];
}
int tm = (tl + tr)/2;
return min(query(v*2, tl, tm, l, min(r, tm)), query(v*2+1, tm+1, tr, max(l, tm+1), r));
}
};
signed main(){
fastio;
int tests;
cin >> tests;
while(tests--){
int n;
cin >> n;
int a[n];
for(int i = 0; i < n; i++){
cin >> a[i];
}
int left[n];
int right[n];
stack<pair<int, int> > s;
for(int i = 0; i < n; i++){
while(!s.empty()){
int t = s.top().first;
if(t > a[i]) break;
s.pop();
}
if(s.empty()){
left[i] = -1;
}
else{
left[i] = s.top().second;
}
s.push(make_pair(a[i], i));
}
while(!s.empty()) s.pop();
for(int i = n-1; i >= 0; i--){
while(!s.empty()){
int t = s.top().first;
if(t >= a[i]) break;
s.pop();
}
if(s.empty()){
right[i] = -1;
}
else{
right[i] = s.top().second;
}
s.push(make_pair(a[i], i));
}
int prefix[n];
for(int i = 0; i < n; i++){
prefix[i] = 0;
if(i>0) prefix[i] = prefix[i-1];
prefix[i] += a[i];
}
int shifted_prefix[n];
for(int i = 0; i < n; i++){
if(i==0){
shifted_prefix[i] = 0;
continue;
}
shifted_prefix[i] = prefix[i-1];
}
SegtreeMax pref_mx(n), get_border(n);
SegtreeMin pref_mi(n);
pref_mx.build(prefix, 1, 0, n-1);
pref_mi.build(shifted_prefix, 1, 0, n-1);
get_border.build(a, 1, 0, n-1);
int ans = LL_MIN;
for(int i = 0; i <= n-1; i++){
int le_mx = left[i];
int ri_mx = right[i];
if(left[i]==-1) le_mx = -1;
if(right[i]==-1) ri_mx = n;
if(left[i]!=-1){
// choose left[i] as 2nd max
int left_border = get_border.query2(1, 0, n-1, 0, le_mx-1, a[i], true);
if(left[i] == 0) left_border = -1;
// can take from left_border+1 to le_mx
int min_prefix = pref_mi.query(1, 0, n-1, left_border+1, le_mx);
int max_prefix = pref_mx.query(1, 0, n-1, i, ri_mx-1);
ans = max(ans, max_prefix - min_prefix - a[i] - a[le_mx]);
}
if(right[i]!=-1){
// chose right[i] as 2nd max
int right_border = get_border.query2(1, 0, n-1, ri_mx+1, n-1, a[i], false);
if(ri_mx == n-1) right_border = n;
if(right_border == -1) right_border = n;
int min_prefix = pref_mi.query(1, 0, n-1, le_mx+1, i);
int max_prefix = pref_mx.query(1, 0, n-1, ri_mx, right_border-1);
ans = max(ans, max_prefix - min_prefix - a[i] - a[ri_mx]);
}
}
cout << ans << endl;
}
return 0;
}
Editorialist's Code
#include "bits/stdc++.h"
// #pragma GCC optimize("Ofast,unroll-loops")
// #pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,mmx,avx,avx2")
using namespace std;
using ll = long long;
const array<ll, 3> unit = {LLONG_MAX, LLONG_MIN, LLONG_MIN};
template<class T>
struct SegTree {
T f(T a, T b) { return {min(a[0], b[0]), max(a[1], b[1]), max(a[2], b[2])}; }
vector<T> s; int n;
SegTree(int _n = 0, T def = unit) : s(2*_n, def), n(_n) {}
void update(int pos, T val) {
for (s[pos += n] = val; pos /= 2;)
s[pos] = f(s[pos * 2], s[pos * 2 + 1]);
}
T query(int b, int e) {
T ra = unit, rb = unit;
if (b == -1) {
ra[0] = rb[0] = 0;
++b;
}
for (b += n, e += n; b < e; b /= 2, e /= 2) {
if (b % 2) ra = f(ra, s[b++]);
if (e % 2) rb = f(s[--e], rb);
}
return f(ra, rb);
}
};
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector<ll> a(n), pref(n);
for (ll &x : a)
cin >> x;
SegTree<array<ll, 3>> T(n);
for (int i = 0; i < n; ++i) {
pref[i] = a[i];
if (i) pref[i] += pref[i-1];
T.update(i, array{pref[i], pref[i], a[i]});
}
// left[i] -> highest j < i such that a[j] > a[i]
// right[i] -> lowest j > i such that a[j] < a[i]
vector<int> left(n, -1), right(n, n);
stack<int> s;
for (int i = 0; i < n; ++i) {
while (!s.empty()) {
if (a[s.top()] < a[i]) s.pop();
else break;
}
if (!s.empty()) left[i] = s.top();
s.push(i);
}
while (!s.empty()) s.pop();
for (int i = n-1; i >= 0; --i) {
while (!s.empty()) {
if (a[s.top()] <= a[i]) s.pop();
else break;
}
if (!s.empty()) right[i] = s.top();
s.push(i);
}
ll ans = 0;
auto find_last = [&] (int R, int x) {
if (R <= -1 or T.query(0, R+1)[2] <= x) return -1;
int L = 0;
while (L < R) {
int mid = (L+R+1)/2;
if (T.query(mid, R+1)[2] > x) L = mid;
else R = mid - 1;
}
return L;
};
auto find_first = [&] (int L, int x) {
if (T.query(L, n)[2] <= x) return n;
int R = n-1;
while (L < R) {
int mid = (L+R)/2;
if (T.query(L, mid+1)[2] > x) R = mid;
else L = mid+1;
}
return L;
};
for (int i = 0; i < n; ++i) {
int L = left[i], R = right[i];
if (L != -1) { // Max on left
ll midsum = pref[i-1] - pref[L];
ll rtsum = T.query(i, R)[1] - pref[i];
ll ltsum = 0;
if (L >= 1) {
int qpos = find_last(L-1, a[i]);
ltsum = pref[L-1] - T.query(qpos, L)[0];
}
ans = max(ans, midsum + ltsum + rtsum);
}
if (R != n) { // Max on right
ll midsum = pref[R-1] - pref[i];
ll ltsum = 0, rtsum = 0;
if (i > 0) {
ltsum = pref[i-1] - T.query(L, i)[0];
}
if (R+1 < n) {
int qpos = find_first(R+1, a[i]);
rtsum = T.query(R, qpos)[1] - pref[R];
}
ans = max(ans, midsum + ltsum + rtsum);
}
}
cout << ans << '\n';
}
}
