kmjp's blog

競技プログラミング参加記です

Codeforces ECR #191 : D. Goods on the Shelf

コーナーケース漏れが怖い。
https://codeforces.com/contest/2233/problem/D

問題

整数列Aが与えられる。
最大1回まで2要素をswapできるとき、同じ値が数列中で連続するようにできるか。

解法

初期状態で、連続していない値の数を求めよう。
そのような数が0個なら初期状態で条件達成、3個以上なら達成不可である。

  • 不連続な値が1個(=x)の場合、以下のパターンがある。
    • 2個の連続領域を1個にまとめる
    • 3個の連続領域のうち、1個を残り2個の間に埋める
  • 不連続な値が2個(=x,y)の場合、それぞれだけに着目して、連続な状態に持っていける交換位置の対をつくる。x,yに対し、同じ交換位置の対があるなら条件を満たす。
int T,N,A[202020],C[202020];
vector<pair<int,int>> P[202020];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		FOR(i,N) {
			P[i].clear();
			C[i]=0;
		}
		cin>>N;
		vector<int> As;
		FOR(i,N) {
			cin>>A[i];
			As.push_back(A[i]);
		}
		sort(ALL(As));
		FOR(i,N) {
			A[i]=lower_bound(ALL(As),A[i])-As.begin();
			if(P[A[i]].empty()||P[A[i]].back().second+1!=i) {
				P[A[i]].push_back({i,i});
			}
			else {
				P[A[i]].back().second=i;
			}
			C[A[i]]++;
		}
		vector<int> F;
		FOR(i,N) {
			if(P[i].size()>=2) F.push_back(i);
			if(P[i].size()>=4) {
				F.push_back(i);
				F.push_back(i);
				F.push_back(i);
			}
		}
		
		if(F.empty()) {
			cout<<"YES"<<endl;
			continue;
		}
		if(F.size()>=3) {
			cout<<"NO"<<endl;
			continue;
		}
		
		int ok=0;
		if(F.size()==1) {
			x=F[0];
			if(P[x].size()==2) {
				if(P[x][0].second+2==P[x][1].first) ok=1;
				if(P[x][0].first==P[x][0].second&&A[P[x][0].second+1]==A[P[x][1].first-1]) ok=1;
				if(P[x][0].first==P[x][0].second&&C[A[P[x][1].first-1]]==1) ok=1;
				if(P[x][0].first==P[x][0].second&&P[x][1].second+1<N&&C[A[P[x][1].second+1]]==1) ok=1;
				if(P[x][1].first==P[x][1].second&&A[P[x][0].second+1]==A[P[x][1].first-1]) ok=1;
				if(P[x][1].first==P[x][1].second&&C[A[P[x][0].second+1]]==1) ok=1;
				if(P[x][1].first==P[x][1].second&&P[x][0].first&&C[A[P[x][0].first-1]]==1) ok=1;
			}
			if(P[x].size()==3) {
				if(P[x][0].first==P[x][0].second&&P[x][1].second+2==P[x][2].first) ok=1;
				if(P[x][2].first==P[x][2].second&&P[x][0].second+2==P[x][1].first) ok=1;
			}
		}
		else {
			x=F[0];
			y=F[1];
			set<pair<int,int>> C;
			if(P[x].size()==2) {
				if(P[x][0].second+2==P[x][1].first) {
					C.insert({P[x][0].first,P[x][0].second+1});
					C.insert({P[x][0].second+1,P[x][1].second});
				}
				if(P[x][0].first==P[x][0].second) {
					C.insert({P[x][0].first,P[x][1].first-1});
					if(P[x][1].second+1<N) C.insert({P[x][0].first,P[x][1].second+1});
				}
				if(P[x][1].first==P[x][1].second) {
					C.insert({P[x][0].second+1,P[x][1].first});
					if(P[x][0].first) C.insert({P[x][0].first-1,P[x][1].first});
				}
			}
			if(P[x].size()==3) {
				if(P[x][0].first==P[x][0].second&&P[x][1].second+2==P[x][2].first) C.insert({P[x][0].first,P[x][2].first-1});
				if(P[x][2].first==P[x][2].second&&P[x][0].second+2==P[x][1].first) C.insert({P[x][0].second+1,P[x][2].first});
			}
			if(P[y].size()==2) {
				if(P[y][0].second+2==P[y][1].first) {
					ok+=C.count({P[y][0].first,P[y][0].second+1});
					ok+=C.count({P[y][0].second+1,P[y][1].second});
				}
				if(P[y][0].first==P[y][0].second) {
					ok+=C.count({P[y][0].first,P[y][1].first-1});
					if(P[y][1].second+1<N) ok+=C.count({P[y][0].first,P[y][1].second+1});
				}
				if(P[y][1].first==P[y][1].second) {
					ok+=C.count({P[y][0].second+1,P[y][1].first});
					if(P[y][0].first) ok+=C.count({P[y][0].first-1,P[y][1].first});
				}
			}
			if(P[y].size()==3) {
				if(P[y][0].first==P[y][0].second&&P[y][1].second+2==P[y][2].first) ok+=C.count({P[y][0].first,P[y][2].first-1});
				if(P[y][2].first==P[y][2].second&&P[y][0].second+2==P[y][1].first) ok+=C.count({P[y][0].second+1,P[y][2].first});
			}
			
			
		}
		if(ok) {
			cout<<"YES"<<endl;
		}
		else {
			cout<<"NO"<<endl;
		}
		
		
		
	}
}

まとめ

場合分けが結構面倒。