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

Easy

PREREQUISITES:

Dynamic Programming

PROBLEM:

An array B is said to be good if, for each index i and element 1 \le x \lt B_i, there exists an index j \lt i such that B_j = x.

Given an array A, count the number of good subarrays.

EXPLANATION:

First, let’s understand when an array is good.

Before each element, there must exist at least one occurrence of every positive integer smaller than it.
Such a condition is hardest to satisfy for the leftmost occurrence of any given element; since if the leftmost occurrence has the condition satisfied, all other occurrences will satisfy it too.

So, let \text{first}[x] denote the index of the leftmost occurrence of x.
Then, it can be seen that an array is good if and only if

\text{first}[1] \lt \text{first}[2] \lt \text{first}[3] \lt\ldots

That is, the first occurrences of the elements must be in ascending order.

We’ll use this observation to solve the problem.

Let’s fix an index L, and try to find all indices R such that subarray A[L, R] is good.

Applying the above condition about first occurrences, it’s easy to see that if A[L, R] is good, then A[L, x] will also be good for all x \le R.
On the other hand, if A[L, R] is not good, neither will be A[L, x] for all x \ge R.
This is quite simply because appending a new element does not change first occurrences of all previous elements.

Hence, with L fixed, it is enough for us to find the largest R such that A[L, R] is good.
Let’s now try to find such an R.

First, if A_L \gt 1, it’s impossible for any subarray starting at L to be good; so we ignore such R entirely.

Now, consider A_L = 1.
Let x \gt L be the first index containing an element that’s \gt 1.

Note that every subarray ending before x will certainly be good (since it’ll only contain 1's).
Once we include index x however,

  • If A_x \gt 2, no further subarray can be good (since we’ve encountered a larger value but not 2.)
  • If A_x = 2 however, some further subarrays might be good.
    • In fact, to find the next “limit”, we can repeat the same process starting from index x instead.
    • That is, let y \gt x be the nearest index containing a value that’s \gt 2.
      Then, everything till index y-1 is good; and depending on whether A_y = 3 or not we either continue on with the same process or stop immediately.

Of course, running this algorithm can take \mathcal{O}(N) time from a single starting index, so running it for every L results in quadratic runtime which is too slow.

To speed it up, observe that starting at L, we’re repeatedly jumping to the next greater element - only stopping exactly when the next greater element is larger than the current element plus one.

So, suppose we define dp_i to be the endpoint of this process assuming we start jumping from index i.
Then,

  • Let j \gt i be the index of the nearest element to the right that’s \gt A_i.
  • If A_j = A_i + 1, then we jump to j and continue jumping.
    So, we have dp_i = dp_j.
  • If A_j \gt A_i + 1, then we must stop right before j.
    So, we have dp_i = j-1 in this case.

Assuming the index j can be found quickly for a fixed i, this allows us to compute all the dp_i values in linear time.

Computing the next greater element for every element of an array in linear time is a standard problem that can be solved using a stack,

Once the dp_i values are known, computing the actual answer is easy.
For each index i such that A_i = 1, the largest valid index is exactly dp_i.
So, there are (dp_i - i + 1) valid endpoints for this starting point.
Sum this up across all valid i to obtain the final answer.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Author's code (C++) 
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18

mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());

void Solve() 
{
    int n; cin >> n;
    
    vector <int> a(n + 1);
    for (int i = 1; i <= n; i++){
        cin >> a[i];
    }
    
    vector <int> nxg(n + 1, n + 1);
    stack <pair<int, int>> st;
    
    for (int i = 1; i <= n; i++){
        while (!st.empty() && a[i] > st.top().first){
            nxg[st.top().second] = i;
            st.pop();
        }
        
        st.push({a[i], i});
    }
    
    vector <int> dp(n + 1);
    for (int i = n; i >= 1; i--){
        if (nxg[i] > n){
            dp[i] = n;
        } else {
            if (a[nxg[i]] == a[i] + 1){
                dp[i] = dp[nxg[i]];
            } else {
                dp[i] = nxg[i] - 1;
            }
        }
    }
    
    int ans = 0;
    for (int i = 1; i <= n; i++) if (a[i] == 1){
        ans += dp[i] - i + 1;
    }
    
    cout << ans << "\n";
}

int32_t main() 
{
    auto begin = std::chrono::high_resolution_clock::now();
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    int t = 1;
    // freopen("in",  "r", stdin);
    // freopen("out", "w", stdout);
    
    cin >> t;
    for(int i = 1; i <= t; i++) 
    {
        //cout << "Case #" << i << ": ";
        Solve();
    }
    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
    cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n"; 
    return 0;
}