yukicoder : No.824 Many Shifts Hard

考察は大変だけど実装は軽いので★4～4.5でもよさそう。
https://yukicoder.me/problems/no/824

問題

0～N番のマスがある。
初期状態で1～Nのマスに1個ずつ駒が置かれている。
1～Nの番号を1つ選び、対応する番号のマスに乗る駒をすべて1つ小さい番号に移す、ということをK回行う。

番号の選び方はN^K通りあるが、それぞれにおいて、最終状態で駒が1個以上乗っているマスの番号の総和の総和を求めよ。

解法

初期状態でt番と(t+1)番にのる駒が途中で合流しない場合、(t+1)番の駒の最終位置のマスの番号分が解に寄与する。
すなわち、初期状態として2つの駒の位置(t,t+1)を取り、K回番号の過程で両者が合流することなく選択後(a,b)に到達する組み合わせがn通りとすると、n*bだけ解に寄与する。

途中先行する(t番のマスから始まる)駒が0番のマスに到達すると、それらはそれ以上動かない。
よって、途中で先行するコマが0番に到達しかねないK番以前の駒は個別に数え上げ、到達しえない(K+1)番以降の駒はまとめて数え上げよう。

int N,K;
ll from[302][302][2];
ll to[302][302][2];
ll mo=1000000007;
ll ret;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	
	for(i=1;i<=min(N,K);i++) from[i-1][i][0]++;
	from[K][K+1][1]=1;
	FOR(i,K) {
		ZERO(to);
		for(x=0;x<=K;x++) for(y=x+1;y<=K+1;y++) FOR(j,2) {
			if(x==0) {
				// take y
				if(y-1>x) (to[x][y-1][j]+=from[x][y][j])%=mo;
				// other
				(to[x][y][j]+=from[x][y][j]*(N-1))%=mo;
			}
			else {
				// take x;
				(to[x-1][y][j]+=from[x][y][j])%=mo;
				// take y
				if(y-1>x) (to[x][y-1][j]+=from[x][y][j])%=mo;
				// other
				(to[x][y][j]+=from[x][y][j]*(N-2))%=mo;
			}
		}
		
		swap(from,to);
	}
	
	for(x=0;x<=K;x++) for(y=x+1;y<=K;y++) (ret+=y*from[x][y][0])%=mo;
	for(y=1;y<=K+1;y++) {
		ll pat=0;
		for(x=0;x<y;x++) pat+=from[x][y][1];
		pat%=mo;
		for(i=K+1;i<=N;i++) (ret+=pat*(i-(K+1-y)))%=mo;
	}
	cout<<ret<<endl;
}

まとめ

「t番と(t+1)番にのる駒が途中で合流しない場合」という数え上げ方に考えが至ると、あとは割とスムーズ。