Codeforces #838 : Div2 F. Good Pairs

コードはそこまで長くない。
https://codeforces.com/contest/1762/problem/F

問題

N要素の整数列Aと、Kが与えられる。
以下を満たす整数対(L,R)はいくつあるか。

• A[L]....A[R]の部分列のうち、先頭はA[L]、末尾はA[R]としたとき、部分列の隣接要素の差の絶対値がK以下となるようにできるか。

解法

A[L]=A[R]の場合は明らか。
A[L]＜A[R]の場合を考えると、A[L]からできるだけ大きくできるように部分列を拾い上げて行くと良い。

部分列を昇順となるように拾うことを考える。
iの小さい順に見て行き、部分列にA[i]を拾える時、次はA[i]+1～A[i]+Kを拾える、という条件をもとに、BITやSegTreeで次に各値を拾えるような左端は何通りあるかを数え上げて行こう。

int T;
int N,K;
int A[505050];

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<int,21> bt;

ll hoge() {
	map<int,int> M;
	int i;
	ll ret=0;
	FOR(i,N) {
		ret+=bt(A[i]);
		bt.add(A[i]+1,1);
		bt.add(A[i]+K+1,-1);
		M[A[i]+K]++;
		while(1) {
			auto it=M.lower_bound(A[i]);
			if(it==M.end()||it->first>=A[i]+K) break;
			bt.add(it->first+1,it->second);
			bt.add(A[i]+K+1,-it->second);
			int v=it->second;
			M.erase(it);
			M[A[i]+K]+=v;
		}
	}
	FOR(i,N) bt.add(A[i]+1,-1);
	FORR2(a,b,M) bt.add(a+1,b);
	return ret;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K;
		map<int,int> S;
		ll ret=0;
		FOR(i,N) {
			cin>>A[i];
			ret+=++S[A[i]];
		}
		ret+=hoge();
		reverse(A,A+N);
		ret+=hoge();
		cout<<ret<<endl;
	}
}

まとめ

これBITで書いたけど、SegTreeの方が楽だったかもな。