CPSCO2019 Session3 : F - Flexible Permutation

O(N^2)にできるかな？と思いながら解いてた。
https://atcoder.jp/contests/cpsco2019-s3/tasks/cpsco2019_s3_f

問題

正整数N,A,Bが与えられる。
1～NのPermutation　Pのうち、P[i]＞iとなるiがA個、P[i]＜iとなるiがB個あるような並び順は何通りか。

解法

挿入DPで解く。
dp(n,a,b) := 1～nの並び順を決めたとき、P[i]＞iとなるiがa個、P[i]＜iとなるiがb個ある並び順
とし、状態遷移を考える。

1～nのPermutationのどこかに(n+1)を加えることを考えよう。

• 末尾に加えると、a,bは変化しない
  • dp(n+1,a,b) += dp(n,a,b)
• P[i]＞iである場所とswapすると、n+1＞iだがP[i]＜n+1なのでbが1個増える
  • dp(n+1,a,b+1) += a*dp(n,a,b)
• P[i]＜iである場所とswapすると、n+1＞iだがP[i]＜n+1なのでaが1個増える
  • dp(n+1,a+1,b) += b*dp(n,a,b)
• P[i]=iである場所とswapすると、n+1＞iだがP[i]＜n+1なのでa,bが1個増える
  • dp(n+1,a+1,b+1) += (n-a-b)*dp(n,a,b)

あとはdp(N,A,B)が解。
上記DPはO(N^3)かかるが、先にP[i]=iである箇所をComb(N,N-A-B)通り固定しておけば、状態数を1個減らしO(N^2)でも解ける。

int N,A,B;
ll dp[303][303][303];
ll mo=1000000007;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>A>>B;
	dp[0][0][0]=1;
	FOR(i,N) {
		FOR(x,A+1) FOR(y,B+1) {
			// last
			(dp[i+1][x][y]+=dp[i][x][y])%=mo;
			// swap a
			(dp[i+1][x][y+1]+=dp[i][x][y]*x)%=mo;
			// swap b
			(dp[i+1][x+1][y]+=dp[i][x][y]*y)%=mo;
			// swap =
			(dp[i+1][x+1][y+1]+=dp[i][x][y]*(i-x-y))%=mo;
		}
	}
	
	
	cout<<dp[N][A][B]<<endl;
}

まとめ

挿入DPに行くまでちょっと手間取った。