Problem p03522 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

#define fi first
#define se second
#define mp make_pair
#define pb push_back

typedef long long ll;
typedef pair<int, int> ii;
typedef vector<int> vi;
typedef long double ld;
typedef tree<ii, null_type, less<ii>, rb_tree_tag,
             tree_order_statistics_node_update>
    pbds;

const int N = 455555;
const int INF = int(1e9);
int a[N];

/*
struct node
{
        int v,lazy;
};
node st[N*4];
void build(int id, int l, int r)
{
        st[id].lazy=0;st[id].v=INF;
        if(r-l>=2)
        {
                int mid=(l+r)>>1;
                build(id*2,l,mid); build(id*2+1,mid,r);
        }
}

void push(int id, int l, int r)
{
        if(st[id].lazy!=0)
        {
                st[id].v+=st[id].lazy;
                if(r-l>=2)
                {
                        st[id*2].lazy+=st[id].lazy;
st[id*2+1].lazy+=st[id].lazy;
                }
                st[id].lazy=0;
        }
}

void update(int id, int l, int r, int pos, int v)
{
        push(id,l,r);
        if(pos>=r||pos<l) return ;
        if(r-l<2)
        {
                st[id].v=v; return ;
        }
        int mid=(l+r)>>1;
        update(id*2,l,mid,pos,v); update(id*2+1,mid,r,pos,v);
        st[id].v=min(st[id*2].v,st[id*2+1].v);
}

void updaterange(int id, int l, int r, int ql, int qr, int v)
{
        push(id,l,r);
        if(ql>=r||l>=qr) return ;
        if(ql<=l&&r<=qr)
        {
                st[id].lazy+=v; push(id,l,r); return ;
        }
        int mid=(l+r)>>1;
        updaterange(id*2,l,mid,ql,qr,v); updaterange(id*2+1,mid,r,ql,qr,v);
        st[id].v=min(st[id*2].v,st[id*2+1].v);
}

ii query(int id, int l, int r, int ql, int qr)
{
        push(id,l,r);
        if(ql>=r||l>=qr) return mp(INF, -1);
        if(ql<=l&&r<=qr) return mp(st[id].v, l);
        int mid=(l+r)>>1;
        return min(query(id*2,l,mid,ql,qr),query(id*2+1,mid,r,ql,qr));
}
*/

struct node {
  int prior, siz;
  ii val;   // value stored in the array
  ii sum;   // whatever info you want to maintain in segtree for each node
  int lazy; // whatever lazy update you want to do
  struct node *l, *r;
};

node buffer[555555];
typedef node *pnode;

int sz(pnode t) { return t ? t->siz : 0; }
void upd_sz(pnode t) {
  if (t)
    t->siz = sz(t->l) + 1 + sz(t->r);
}
void lazy(pnode t) {
  if (!t || !t->lazy)
    return;
  t->val.fi += t->lazy; // operation of lazy
  t->sum.fi += t->lazy;
  if (t->l)
    t->l->lazy += t->lazy; // propagate lazy
  if (t->r)
    t->r->lazy += t->lazy;
  t->lazy = 0;
}
void reset(pnode t) {
  if (t)
    t->sum = t->val; // no need to reset lazy coz when we call this lazy would
                     // itself be propagated
}
void combine(pnode &t, pnode l, pnode r) { // combining two ranges of segtree
  if (!l || !r)
    return void(t = l ? l : r);
  t->sum = min(l->sum, r->sum);
}
void operation(pnode t) { // operation of segtree
  if (!t)
    return;
  reset(t); // reset the value of current node assuming it now represents a
            // single element of the array
  lazy(t->l);
  lazy(t->r); // imp:propagate lazy before combining t->l,t->r;
  combine(t, t->l, t);
  combine(t, t, t->r);
}
void split(pnode t, pnode &l, pnode &r, int pos, int add = 0) {
  // if(pos<0&&add==0){l=NULL; r=t; return ;}
  if (!t)
    return void(l = r = NULL);
  lazy(t);
  int curr_pos = add + sz(t->l);
  if (curr_pos <= pos) // element at pos goes to left subtree(l)
    split(t->r, t->r, r, pos, curr_pos + 1), l = t;
  else
    split(t->l, l, t->l, pos, add), r = t;
  upd_sz(t);
  operation(t);
}
void merge(pnode &t, pnode l,
           pnode r) { // l->leftarray,r->rightarray,t->resulting array
  lazy(l);
  lazy(r);
  if (!l || !r)
    t = l ? l : r;
  else if (l->prior > r->prior)
    merge(l->r, l->r, r), t = l;
  else
    merge(r->l, l, r->l), t = r;
  upd_sz(t);
  operation(t);
}

pnode curnode = &buffer[0];
pnode init(ii val) {
  pnode ret = curnode++;
  ret->prior = rand();
  ret->siz = 1;
  ret->val = val;
  ret->sum = val;
  ret->lazy = 0;
  return ret;
}
ii range_query(pnode t, int l, int r) { //[l,r]
  if (l > r)
    return mp(INF, -1);
  pnode L, mid, R;
  split(t, L, mid, l - 1);
  split(mid, t, R, r - l); // note: r-l!!
  ii ans = t->sum;
  merge(mid, L, t);
  merge(t, mid, R);
  return ans;
}
void range_update(pnode t, int l, int r, ll val) { //[l,r]
  if (l > r)
    return;
  pnode L, mid, R;
  split(t, L, mid, l - 1);
  split(mid, t, R, r - l); // note: r-l!!
  t->lazy += val;          // lazy_update
  merge(mid, L, t);
  merge(t, mid, R);
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  int n;
  cin >> n;
  for (int i = 0; i < n; i++)
    cin >> a[i];
  pbds T;
  int timer = 0;
  ll ans = 0;
  pnode treap = NULL;
  for (int i = 0; i < n; i++) {
    int aft = T.order_of_key(mp(a[i] + 1, -1)) - 1;
    int mn = i - 1 - aft;
    if (mn == 0) {
      T.insert(mp(a[i], ++timer));
      pnode tmp = init(mp(a[i] - i, a[i]));
      merge(treap, treap, tmp);
      continue;
    }
    int id = aft + 1;
    aft = -1;
    ii X = range_query(treap, id + 1, i - 1);
    if (X.fi + i - 1 - a[i] < mn) {
      mn = X.fi + i - 1 - a[i];
      aft = X.se;
    }
    ans += mn;
    // cerr<<"MN : "<<mn<<" AFT : "<<aft<<endl;
    if (aft >= 0) {
      int pos = T.order_of_key(mp(aft + 1, -1));
      range_update(treap, pos, i - 1, -1);
      T.insert(mp(aft, ++timer));
      pnode tmp = init(mp(aft - pos, aft));
      pnode L, R;
      split(treap, L, R, pos - 1);
      merge(L, L, tmp);
      merge(treap, L, R);
    } else {
      aft = a[i];
      int pos = T.order_of_key(mp(aft + 1, -1));
      range_update(treap, pos, i - 1, -1);
      T.insert(mp(aft, ++timer));
      pnode tmp = init(mp(aft - pos, aft));
      pnode L, R;
      split(treap, L, R, pos - 1);
      merge(L, L, tmp);
      merge(treap, L, R);
    }
    // cerr<<"MN "<<mn<<' '<<"AFT "<<aft<<'\\n';
  }
  cout << ans << '\\n';
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

#define fi first
#define se second
#define mp make_pair
#define pb push_back

typedef long long ll;
typedef pair<int, int> ii;
typedef vector<int> vi;
typedef long double ld;
typedef tree<ii, null_type, less<ii>, rb_tree_tag,
             tree_order_statistics_node_update>
    pbds;

int a[222222];

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  int n;
  cin >> n;
  for (int i = 0; i < n; i++)
    cin >> a[i];
  vector<int> vec;
  ll ans = 0;
  for (int i = 0; i < n; i++) {
    if (vec.empty() || a[i] >= vec.back()) {
      vec.pb(a[i]);
      continue;
    }
    int mn = int(1e9);
    int aft = i - 1;
    int x = 0;
    for (int j = i - 1; j >= -1; j--) {
      int val = i - 1 - j + max(0, (j >= 0 ? vec[j] : -int(1e9)) - a[i]);
      if (val < mn) {
        mn = val;
        aft = j;
        x = max(0, (j >= 0 ? vec[j] : -int(1e9)) - a[i]);
      }
    }
    vi nw;
    ans += mn;
    for (int j = 0; j <= aft; j++)
      nw.pb(vec[j]);
    nw.pb(a[i] + x);
    for (int j = aft + 1; j < i; j++)
      nw.pb(vec[j]);
    vec = nw;
  }
  cout << ans << '\\n';
}