kmjp's blog

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

AtCoder ARC #100 : F - Colorful Sequences

これ系苦手…。
https://beta.atcoder.jp/contests/arc100/tasks/arc100_d

問題

整数N,Kおよび長さMの数列Aが与えられる。
各要素1~Kで構成される長さNの数列を考える。

この数列がカラフルであるとは、長さKの部分列には必ず1~Kが1個ずつ含まれることを示す。
全てのカラフルな数列のうち、Aの登場回数の総和を求めよ。

解法

Aがカラフルでない場合解は0。
以下Aがカラフルであることを前提とする。

今回の例では、1つの数列中にAが複数回あれば複数回その数列がカウントされる。
よってAの登場位置を固定したときその前後の要素について考える、ということを各位置について行えばよい。

Aの位置を固定すると、元の数列がカラフルかどうか問わないならばその組み合わせはK^(N-M)である。
ここから、カラフルでないケースを引くことを考えよう。

まずAは忘れてある長さの数列のうちカラフルでないものの組み合わせを考えよう。
dp(a,b):= 長さaの数列で、末尾b個の要素が異なっており、末尾(b+1)個の要素までを考えると重複がある場合の組み合わせ
dp(a,b)からdp(a+1,b')への状態遷移は、末尾b個の要素を追加するかそれ以外の要素を追加するの2通りなので、dp(N,K)のテーブルはO(NK)で埋めることができる。
実際には、さらにこれに「途中でカラフルじゃない状態になったことがある・ない」も含めたテーブルを作っておく。

  • A中に重複が1個もない場合、カラフルじゃない数列は、先のDPで求めたすべてのカラフルじゃない数列のうちK!/(K!-M!)で割った数
  • A中に重複がある場合、Aのprefixとsuffixでそれぞれいくつ目まで重複がないかを先求めておき、その値を元に先のテーブルを用いてAの前後に要素を追加したときにカラフルでないものが出来上がる数を求める。
int N,K,M;
int A[252525];
int cnt[404];
ll mo=1000000007;

ll dp[25205][2][404], dp2[2][404];

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K>>M;
	FOR(i,M) cin>>A[i], A[i]--;
	
	ll tot=(N-M+1);
	FOR(i,N-M) (tot*=K)%=mo;
	
	for(i=0;i+K<=M;i++) {
		ZERO(cnt);
		int ok=0;
		for(j=i;j<i+K;j++) if(++cnt[A[j]]==2) ok++;
		if(ok==0) {
			cout<<tot<<endl;
			return;
		}
	}
	
	
	dp[1][0][1]=K;
	dp[1][M==1][1]=K;
	for(i=2;i<=N;i++) {
		ZERO(dp2);
		for(j=1;j<K;j++) {
			(dp[i][0][j+1]+=(K-j)*dp[i-1][0][j])%=mo;
			(dp[i][1][j+1]+=(K-j)*dp[i-1][1][j])%=mo;
			dp2[0][1]+=dp[i-1][0][j];
			dp2[0][j+1]+=mo-dp[i-1][0][j];
			dp2[1][1]+=dp[i-1][1][j];
			dp2[1][j+1]+=mo-dp[i-1][1][j];
		}
		
		for(j=1;j<K;j++) {
			dp2[0][j]+=dp2[0][j-1];
			(dp[i][0][j]+=dp2[0][j])%=mo;
			dp2[1][j]+=dp2[1][j-1];
			(dp[i][1][j]+=dp2[1][j])%=mo;
			if(j>=M) (dp[i][1][j]+=dp[i][0][j])%=mo;
		}
	}
	ZERO(cnt);
	FOR(i,M) cnt[A[i]]++;
	FOR(i,K) if(cnt[i]>1) break;
	
	ll ret=0;
	if(i<K) {
		int F,B;
		set<int> S;
		FOR(F,M) {
			if(S.count(A[F])) break;
			S.insert(A[F]);
		}
		S.clear();
		FOR(B,M) {
			if(S.count(A[M-1-B])) break;
			S.insert(A[M-1-B]);
		}
		for(int L=0;L<=N-M;L++) {
			int R=N-M-L;
			ll LP=0,RP=0;
			for(j=F;j<K;j++) LP+=dp[L+F][0][j];
			for(j=B;j<K;j++) RP+=dp[R+B][0][j];
			LP%=mo;
			RP%=mo;
			ll LM=1,RM=1;
			for(i=K-F+1;i<=K;i++) LM=LM*i%mo;
			for(i=K-B+1;i<=K;i++) RM=RM*i%mo;
			(ret += (LP*modpow(LM))%mo*(RP*modpow(RM)%mo))%=mo;
		}
		
	}
	else {
		// no same
		for(j=1;j<K;j++) ret+=dp[N][1][j];
		ret%=mo;
		ll di=1;
		for(i=K-M+1;i<=K;i++) di=di*i%mo;
		ret=ret*modpow(di)%mo;
		
		
		
	}
		cout<<(tot-ret+mo)%mo<<endl;
	
}

まとめ

これ考察も重いし1400ptぐらいでもいい気がする…。