kmjp's blog

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

AtCoder ARC #089 : F - ColoringBalls

ARC

なぜロシアのコンテストサイトは数列を誕生日プレゼントにしたがるし、AtCoderは赤青のボールを並べたがるのか。
https://beta.atcoder.jp/contests/arc089/tasks/arc089_d

問題

N個の白いボールが並んでいる。
ここでK個のクエリが順に与えられる。

各クエリは赤青のいずれかである。
クエリに応じ、ボール列の一部区間をその色に上塗りする(区間は空でもよい)。
ただし、青を塗るときはすでに下が赤か青で塗られていなければならない。
全クエリを完了後、ボール列の色の塗り分け方は何通りか。

解法

白ボールで区切られた1つの赤青ボール列を考える。
さらに、連続した同じ色のボールを1つにまとめたものをセットにして考える。
青色のボールが0個のケースは、赤色ボール1個の1通りだが、それ以上青色のボールがB個のケースは、赤色のボール数は両サイドの赤ボールの有無により(B-1)、B、B、(B+1)個の4通りある。

ここで、青色ボールがB個のケースは、いずれも同じクエリパターンから生成できる。
クエリの部分列がrb????と?がQ個続くとき、(Q+1)個Bを含むケースが生成できる。
青色ボールが0個のケースはrから生成できる。

L文字のクエリから生成できるW文字のパターンをf(L,W)とする。
f(L,W)は赤青のボール数4通りについて組み合わせ計算することで容易に計算できる。

さて、N個のボールを白色ボールで区切った際、それぞれ何文字のクエリから生成できるかを降順に並べたものを考える。
例えば[3,2,2,1]は["rb?","rb","rb","r"]から生成できるボール列群に相当する。
この[3,2,2,1]のような数列の組み合わせは、50万通りもないのでDFSで総当たりすればよい。

あとはこの数列[3,2,2,1]について、クエリ文字列からその数列を再現できるか、またN個のボールをそのように塗り分ける組み合わせが何通りかをそれぞれ判定しよう。
後者は先に求めたf(L,W)を用いてDPを求めればよい。

問題は前者である。
ただ、これは[3,2,2,1]についてまず"rb"の部分を先頭から貪欲に取り、残った文字があれば(r,bを取ったより後ろで)適宜配分していくとよい。
例えばクエリ文字列が"rrbrbbrrb"だったら、A～Dを[3,2,2,1]の1～4文字目相当と表すと"ABACBCDA_"のように割り当てる。

意外にTLEが厳しいので、適度に枝刈りをするとよい。

ll comb(int P_,int Q_) {
	static const int N_=1020;
	static ll C_[N_][N_];
	
	if(C_[0][0]==0) {
		int i,j;
		FOR(i,N_) C_[i][0]=C_[i][i]=1;
		for(i=1;i<N_;i++) for(j=1;j<i;j++) C_[i][j]=(C_[i-1][j-1]+C_[i-1][j])%mo;
	}
	if(P_<0 || P_>N_ || Q_<0 || Q_>P_) return 0;
	return C_[P_][Q_];
}
ll hcomb(int P_,int Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);};

int N,K;
string S;
ll mo=1000000007;
ll group[71][71];
int ML[71];
ll fact[72],factr[72],inv[72];
ll tot;
vector<int> G;

void hoge(vector<int> V,vector<ll> C) {
	vector<int> O=V;
	string T=S;
	int TN=T.size(),VN=V.size();
	vector<int> R(VN,-1);
	
	int x=0,y=0,i;
	FOR(i,TN) if(T[i]=='r' && x<VN) {
		R[x]=i;
		V[x]--;
		x++;
		T[i]='_';
		
	}
	if(x!=VN) return;
	x=0;
	FOR(i,TN) if(T[i]=='b') {
		if(x<VN && V[x] && i>R[x]) {
			V[x]--;
			R[x]=i;
			x++;
			T[i]='_';
		}
	}
	if(x<VN && V[x]) return;
	
	x=0;
	FOR(i,TN) if(T[i]!='_') {
		if(x<VN && V[x] && i>R[x]) {
			V[x]--;
			T[i]='_';
			if(V[x]==0) x++;
		}
	}
	if(x<VN && V[x]) return;
	
	map<int,int> M;
	FORR(c,O) M[c]++;
	
	ll ret=0;
	FOR(x,N+1) {
		int left=N-x-(V.size()-1);
		if(left<0) break;
		(ret+=C[x]*hcomb(V.size()+1,left))%=mo;
	}
	
	(ret*=fact[V.size()])%=mo;
	FORR(m,M) (ret*=factr[m.second])%=mo;
	
	(tot+=ret)%=mo;
}

void dfs(int L,int cur,vector<ll> C) {
	if(cur==0) return;
	// not take cur
	dfs(L,cur-1,C);
	
	if(L>=ML[cur]) {
		G.push_back(cur);
		vector<ll> C2(71);
		for(int x=0;x<=70;x++) if(C[x]) {
			for(int y=ML[cur];x+y<=70;y++) {
				(C2[x+y]+=C[x]*group[cur][y])%=mo;
			}
		}
		
		hoge(G,C2);
		dfs(L-(ML[cur]+1),cur,C2);
		G.pop_back();
	}
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K>>S;
	
	inv[1]=fact[0]=factr[0]=1;
	for (int i=2;i<=70;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
	for (int i=1;i<=70;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	
	for(i=1;i<=70;i++) group[1][i]=1;
	ML[1]=1;
	for(i=2;i<=70;i++) {
		vector<int> cand={(i-1)*2-1,(i-1)*2,(i-1)*2,(i-1)*2+1};
		ML[i]=(i-1)*2-1;
		FORR(l,cand) {
			for(x=l;x<=70;x++) (group[i][x]+=hcomb(l,x-l))%=mo;
		}
	}
	
	
	tot=1;
	vector<ll> C(71);
	C[0]=1;
	dfs(N,70,C);
	cout<<tot<<endl;
	
}

まとめ

1100ptよりもう少し高くてよさそう。