UnKoder #04 : Dual Sort

こういうシンプルな問題設定はいいね。
https://www.hackerrank.com/contests/unkoder-04/challenges/dual-sort

問題

1～Nのpermutationである数列A[i]が与えられる。

数値Kが与えられるので、A[i]のうちK個を青、残り(N-K)個を赤に分類したとする。
同じ色の数字同士を任意にスワップすることができる。

適切な色分けを行った場合、得られる辞書順最小のAを求めよ。

解法

辞書順最小を求める定番テクとして、数列の前の方から順にループしてA[i]を求めていく。
その際、A[i]としてあり得る最小の値を以下の要領で求めていく。

• jがA[k] (k＜i)で使用済みの場合、当然jはA[i]として使用不可。
• jがrev[j]番目にある(A[rev[j]]=j)とする。
  • Union-Findでiとrev[j]を暫定的にuniteする。この時、Union-FindによるN個の連結成分群をK個と(N-K)個の組みに分けられるかDPで判定し、分けられないならjはA[i]としてスワップできない。
  • 判定に成功したなら、Union-Findを実際にuniteし、A[i]=jとなるスワップ方法があると判断する。

int N,K;
vector<int> A,R,rev;
int used[51];

template<int um> class UF {
	public:
	vector<int> par,rank,cnt;
	UF() {rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par.push_back(i);}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};

bool enableset(UF<50>& uf) {
	int dp[51]={};
	int i,x;
	dp[0]=1;
	FOR(i,N) if(uf[i]==i) {
		int c=uf.count(i);
		for(x=50;x>=c;x--) dp[x] |= dp[x-c];
	}
	return dp[K]==1;
}

vector<int> dodo(UF<50> uf) {
	int i,x;
	vector<int> d=A;
	int B[51];
	FOR(i,N) {
		vector<int> e,f;
		FOR(x,N) if(uf[x]==i) e.push_back(A[x]),f.push_back(x);
		sort(ALL(e));
		FOR(x,e.size()) d[f[x]]=e[x];
	}
	return d;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	rev.resize(N);
	FOR(i,N) cin>>x, A.push_back(x-1), rev[x-1]=i;
	R=A;
	
	UF<50> uf;
	FOR(i,N) {
		FOR(j,N) if(used[j]==false) {
			UF<50> uf2=uf;
			uf2(i,rev[j]);
			if(enableset(uf2)==false) continue;
			
			used[j]=true;
			uf=uf2;
			A=dodo(uf);
			FOR(x,A.size()) rev[A[x]]=x;
			break;
		}
	}
	
	FOR(i,N) _P("%d%c",A[i]+1,(i==N-1)?'\n':' ');
}

まとめ

まさかここでUnion-Findとはね。