APDIS - Editorial

PROBLEM LINK:

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

Author: iceknight1093
Tester: raysh07
Editorialist: iceknight1093

DIFFICULTY:

Easy

PREREQUISITES:

Dynamic programming

PROBLEM:

An array is called good if it can be made sorted by deleting some elements at indices that form an arithmetic progression with common difference \ge 2.

Given an array, count the number of its subarrays that are good.
N \le 4000

EXPLANATION:

Observe that if the subarray [L, R] is good, then [L, R-1] will also be good (by deleting the exact same set of indices.)
This means we only need to find, for each left endpoint, the farthest right endpoint that results in a good array.

However, it’s not immediately obvious how to tell if a subarray is good or not, or how to maintain this information while changing borders.

Instead, we use a different approach.

Suppose we fix d, the common difference of the deleted indices.
Let’s try to count all subarrays that are good with respect to this common difference.

Define dp_L to be the maximum index R such that A[L, R] is good with respect to d.
We’ll try to compute dp_L for all L.


For a fixed left endpoint L, let i \ge L be the smallest index such that A_i \gt A_{i+1}.
If no such index exists then the suffix starting at L is already sorted, so we have dp_L = N; we thus only care about the case where i exists.

Everything ending at or before i is sorted, so we start out with dp_L = i.
Now, for everything else, certainly either i or i+1 must be deleted; otherwise the subarray won’t be sorted.

Suppose we decide to delete the element at index i.
Then,

  • If A_{i-1} \gt A_{i+1}, sortedness can’t be continued anyway.
    Thus we are limited to a right endpoint of i here.
  • If A_{i-1} \le A_{i+1}, sortedness can continue on.
    Here though, observe that the indices that need to be deleted are already fixed - we’ll remove exactly indices i, i+d, i+2d, \ldots
    So, we only need to figure out the longest we can do this while maintaining sortedness; which you can notice doesn’t really depend on the left endpoint at all.

Taking inspiration from the last point, let’s define r_j to be the largest right endpoint such that the subarray A[j, r_j] will become sorted if the elements at indices j-1+d, j-1+2d, \ldots are removed from it.

Note that if we know all the r_j values, then in the case of deleting index i the largest valid right endpoint is simply r_{i+1} (of course, only when A_{i-1} \le A_{i+1}).

Similarly, for the case where we instead delete index i+1, we’ll want to look at r_{i+2} instead (when it’s valid to do so.)
Take the best of both cases to compute dp_L.

Luckily, computing r_j isn’t too hard.

  • If the subarray from index j to index j+d-2 is not sorted, then r_j equals the largest sorted prefix starting at j.
  • Otherwise, r_j equals either j+d-1, or r_{j+d}, depending on whether A_{j+d-2} \le A_{j+d}.

This way, all the r_j values can be computed in linear time once d is fixed.
Following that, all the dp_L values can also be computed in linear time for this fixed d.


We now know which subarrays are good with respect to a certain d - in particular we know the longest d-good subarray starting at each point.

Since only 2 \le d \le N matter, we can simply try all values of d to obtain all possible information in \mathcal{O}(N^2) time.
For each left endpoint L, just take the farthest right endpoint that’s valid for some value d (i.e the maximum value of dp_L across all d); this is the longest valid subarray starting at L, so add the count of such subarrays to the answer.

TIME COMPLEXITY:

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

CODE:

Tester'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];
    }
    
    // if a[j] > a[j + 1], not deleting either j or j + 1 is bad 
    // if a[j] > a[j + 2], deleting j + 1 is bad 
    
    vector <int> ri(n + 1, 1);
    
    for (int d = 2; d <= n; d++){
        vector <int> ok(d), tt(d);
        int glob = n;
        int ti = 1;
        vector <int> changes;
        
        auto get = [&](int x){
            if (tt[x] == ti){
                return ok[x];
            }
            return glob;
        };
        
        auto upd = [&](int x, int v){
            if (tt[x] != ti) changes.push_back(x);
            ok[x] = v;
            tt[x] = ti;
        };
        
        for (int i = n; i >= 1; i--){
            if (i + 1 <= n && a[i] > a[i + 1]){
                int x1 = i % d, x2 = (i + 1) % d;
                int v1 = get(x1), v2 = get(x2);
                ti++;
                glob = i;
                changes.clear();
                upd(x1, v1);
                upd(x2, v2);
            }
            
            if (i + 2 <= n && a[i] > a[i + 2]){
                upd((i + 1) % d, i + 1);
            }
            
            // find max 
            int mx = -1;
            if (changes.size() < d) mx = glob;
            for (auto x : changes) mx = max(mx, ok[x]);
            ri[i] = max(ri[i], mx);
        }
    }
    
    int ans = 0;
    for (int i = 1; i <= n; i++){
        ans += ri[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;
}