1536D - Omkar and Medians - CodeForces Solution

#include <bits/stdc++.h>

using namespace std;

class SegmentTree {
public:
  SegmentTree(int _n) : n(_n) {
    sum.assign(4*n, 0);
    lazy.assign(4*n, false);
    lazy_update.assign(4*n, 0);
  }

  void increment(int i, int delta){
    range_update(i, i, range_sum(i, i) + delta);
  }

  int range_sum(int i, int j){
    return range_sum(1, 0, n-1, i, j);
  }

  void range_update(int i, int j, int newval){
    range_update(1, 0, n-1, i, j, newval);
  }

private:
  const int n;
  vector<int> sum, lazy_update;
  vector<bool> lazy;

  static int left(int p){
    return 2*p;
  }

  static int right(int p){
    return 2*p + 1;
  }

  void propagate(int p, int L, int R){
    if(!lazy[p]) return;
    sum[p] = (R - L + 1) * lazy_update[p];

    if(L != R){
      lazy[left(p)] = lazy[right(p)] = true;
      lazy_update[left(p)] = lazy_update[right(p)] = lazy_update[p];
    }

    lazy[p] = false;
    lazy_update[p] = 0;
  }

  int range_sum(int p, int L, int R, int i, int j){
    propagate(p, L, R);
    if(L > j || R < i) return 0;
    if(L >= i && R <= j) return sum[p];

    int mid = (L + R) / 2;
    return range_sum(left(p), L, mid, i, j) +
      range_sum(right(p), mid+1, R, i, j);
  }

  void range_update(int p, int L, int R, int i, int j, int newval){
    propagate(p, L, R);
    if(L > j || R < i) return;

    if(L >= i && R <= j){
      lazy[p] = true;
      lazy_update[p] = newval;
      propagate(p, L, R);
      return;
    }

    int mid = (L + R) / 2;
    range_update(left(p), L, mid, i, j, newval);
    range_update(right(p), mid+1, R, i, j, newval);

    sum[p] = sum[left(p)] + sum[right(p)];
  }
};

int main(){
  int t;
  scanf("%d", &t);

  while(t--){
    int n;
    scanf("%d", &n);

    int b[n+1];
    map<int,int> cvt;
    for(int i = 1; i <= n; i++){
      scanf("%d", &b[i]);
      cvt[b[i]] = 0;
    }

    int X = 1;
    for(auto &kv : cvt){
      kv.second = ++X;
    }
    ++X;

    for(int i = 1; i <= n; i++){
      b[i] = cvt[b[i]];
    }

    int a[2*n];
    a[1] = b[1];
    list<int> L = {1, a[1], X};
    list<int>::iterator med_it = L.begin(); med_it++;

    SegmentTree st_le(X+1), st_ge(X+1);
    bool test = true;

    for(int i = 2; i <= n; i++){
      if(b[i] == *med_it){
	st_le.increment(b[i], 1);
	st_ge.increment(b[i], 1);

      } else if(b[i] < *med_it){
	auto prev_it = med_it; prev_it--;

	if(*prev_it > b[i]){
	  test = false;
	  break;
	}

	int tmp = st_le.range_sum(b[i], *med_it) + 2;
	st_le.range_update(b[i], *med_it, 0);
	st_le.increment(b[i], tmp);

	if(*prev_it == b[i]){
	  med_it = prev_it;
	} else {
	  st_le.increment(b[i], -1);
	  med_it = L.insert(med_it, b[i]);
	}
      } else {
	auto next_it = med_it; next_it++;

	if(*next_it < b[i]){
	  test = false;
	  break;
	}

	int tmp = st_ge.range_sum(*med_it, b[i]) + 2;
	st_ge.range_update(*med_it, b[i], 0);
	st_ge.increment(b[i], tmp);

	if(*next_it == b[i]){
	  med_it = next_it;
	} else {
	  st_ge.increment(b[i], -1);
	  med_it = L.insert(next_it, b[i]);
	}
      }
    }

    printf("%s\n", test ? "YES" : "NO");
  }
}