CSAcademy Round #10: D.Subarray Medians

CSA10D.cpp

#include<iostream>
#include<vector>

using namespace std;
#define int long long
#define repeat(i,n) for(int i=0;i<(n);i++)

struct Median{
  int K, k, md;
  vector<int> prv, nxt;
  Median(int N){
    K=k=N;
    md=(N+1)/2-1;
    prv.resize(N); nxt.resize(N);
    repeat(i, N){
      prv[i]=i-1;
      nxt[i]=i+1;
    }
  }
  
  void dl(int v){
    if(v<md){
      if(k%2==0) md=nxt[md];
    }else if(v==md){
      if(k%2==0) md=nxt[md];
      else md=prv[md];
    }else{// v>md
      if(k%2==1) md=prv[md];
    }
    if(nxt[v]!=K) prv[nxt[v]]=prv[v];
    if(prv[v]!=-1) nxt[prv[v]]=nxt[v];
    k--;
  }
  
};

signed main(){
  
  int N; cin>> N;
  vector<int> a(N);
  repeat(i, N) cin>> a[i], a[i]--;
  
  Median foo(N);
  int s=0;
  repeat(i, N){
    Median bar=foo;
    for(int j=N-1; j>=i; j--){
      if((j-i)%2==0) s+=(bar.md+1)*(i+1)*(j+1);
      bar.dl(a[j]);
    }
    foo.dl(a[i]);
  }
  
  cout<< s<< endl;
  
  return 0;
}