Codeforces ECR #022 : E. Army Creation

本番方針は立ったもののそういうデータ構造を知らずに撃沈。
http://codeforces.com/contest/813/problem/E

問題

N要素の数列Aが与えられる。
以下のクエリQ個に答えよ。

• 要素の範囲[L,R]が与えられる。A[L,R]からいくつかの要素を選びたい。ただし同じ値の要素はK個までしか選択できない。
• 最大何個の要素を選べるか。

解法

A[x]と同じ要素のうち、K個手前に出てくる要素の添え字をf(x)する。
すなわちA[f(x)]=A[x]で、かつf(x)＜y＜xにおいてA[y]=A[x]となるものは(K-1)個しかない。
またそのような添え字がない場合はf(x)=-infとする。

こうすると、x番目の要素が選択されるのは同時にf(x)番目の要素が選択されない場合である。
つまりL≦x≦Rかつf(x)＜Lとなるxを数え上げればいいことになる。
これに答えるには、区間内においてある数値以下の要素数を数え上げるデータ構造があればよい。

自分はそのようなものを知らなかったので、この機会に作った。
以下は初期段階で配列の要素が確定済みの場合に利用できる。
SegTreeの要領で、各ノードにはSubTree内の全要素を昇順ソートしたものを持たせよう。

クエリの際はlog(N)個のセグメントにクエリを分割できるので、それぞれで昇順ソートしたリストをlower_boundで数え上げればO(log^2N)でクエリに応えられる。
普通のSegTreeより若干重く、メモリおよび構築がO(NlogN)かかる点に注意。

template<class V,int NV> class SegTree_count {
public:
	vector<V> val[NV*2];
	V comp(V l,V r){ return max(l,r);};
	SegTree_count(){ for(int i=NV;i<NV*2;i++) val[i].push_back(1<<30);}
	
	int getval(int x,int y,V v,int l=0,int r=NV,int k=1) {
		if(r<=x || y<=l) return 0;
		if(x<=l && r<=y) return lower_bound(ALL(val[k]),v+1)-val[k].begin();
		return getval(x,y,v,l,(l+r)/2,k*2)+getval(x,y,v,(l+r)/2,r,k*2+1);
	}
	void build() {
		for(int e=NV-1;e>=1;e--) {
			val[e].clear();
			merge(ALL(val[e*2]),ALL(val[e*2+1]),back_inserter(val[e]));
		}
	}
	void update(int entry, V v) {
		val[NV+entry][0]=v;
	}
};


int N,K;
int A[101010];
SegTree_count<int,1<<20> st;

vector<int> V[101010];
int pre[101010];
int Q;
int L,R;


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		cin>>A[i];
		if(V[A[i]].size()>=K) pre[i]=V[A[i]][V[A[i]].size()-K];
		V[A[i]].push_back(i+1);
		st.update(i+1,pre[i]);
	}
	st.build();
	
	cin>>Q;
	int ret=0;
	while(Q--) {
		cin>>L>>R;
		L=(L+ret)%N;
		R=(R+ret)%N;
		if(L>R) swap(L,R);
		ret=st.getval(L+1,R+2,L);
		cout<<ret<<endl;
	}
}

まとめ

新たなデータ構造の勉強になりました。