yukicoder : No.121 傾向と対策 : 門松列（その２）

よい頭の体操になりました。
http://yukicoder.me/problems/292

問題

N個の門松が１列に並んでいる。
３つ並んだ門松が門松列を成すとは、３つの高さがすべて異なり、真ん中が最高または最低のものである。
それぞれの門松の長さをA[i]とする。

A[i],A[j],A[k]が門松列を成すような(i＜j＜k)のトリオの組み合わせを求めよ。

解法

最初に門松の高さを座標圧縮しておく。

まず「すべてが異なる」を無視して数える。
真ん中の門松となるjを総当たりし、以下の総和を取ればよい。

• 真ん中が最高となるケース、すなわちA[i]＜A[j]となるiの数と、A[j]＞A[k]となるkの数の積
• 真ん中が最低となるケース、すなわちA[i]＞A[j]となるiの数と、A[j]＜A[k]となるkの数の積

これはFenwickTreeを２つ用いることで条件を満たすiやkの数を高速に求めることができる。

ここから、A[i]=A[k]となるケースを引けばよい。
A[0]～A[j-1]の間にある高さhの門松の数をLj(h)、A[j+1]～A[N-1]の間にある高さhの門松の数をRj(h)とする。
各jにおいてA[i]=A[k]となるケースの数S(j)はとなる。
S(j+1)の値はS(j)にL[A[j]]・R[A[j]]の差分を加えるだけで求められる。

int N;
map<int,int> M;
int A[1010000];
int L[1010000],R[1010000];
ll ret,dif;


template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	int lower_bound(V val) {
		V tv=0; int i,ent=0;
		for(i=ME-1;i>=0;i--) if(tv+bit[ent+(1<<i)-1]<val) tv+=bit[ent+(1<<i)-1],ent+=(1<<i);
		return ent;
	}
	V total(int e) {V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void update(int e, V val) {e++; while(e<=1<<ME) bit[e-1]+=val,e+=e&-e;}
};
BIT<int,20> bt,bt2;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i], M[A[i]]=0;
	
	x=0;
	ITR(it,M) it->second=++x;
	FOR(i,N) A[i]=M[A[i]], R[A[i]]++;
	
	FOR(i,N) bt2.update(A[i],1);
	FOR(i,N) {
		bt2.update(A[i],-1);
		ret += 1LL*bt.total(A[i]-1)*bt2.total(A[i]-1);
		ret += 1LL*(i-bt.total(A[i]))*((N-i-1)-bt2.total(A[i]));
		bt.update(A[i],1);
	}
	FOR(i,N) {
		dif -= 1LL*L[A[i]]*R[A[i]];
		ret -= dif;
		L[A[i]]++;
		R[A[i]]--;
		dif += 1LL*L[A[i]]*R[A[i]];
	}
	
	
	cout<<ret-dif<<endl;
}

まとめ

左右が等しいケースをどう省くか、ちょっと悩んだ。