kmjp's blog

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

TopCoder SRM 673 Div1 Medium BearPermutations

SRM

600ptだけど何とか解けて良かった。
https://community.topcoder.com/stat?c=problem_statement&pm=14080

問題

数列に関して、スコアを定義する。
まずこの数列からCartesian tree（デカルト木？適切な日本語がわからない）を生成する。
子を２つ持つ各頂点に対し、元数列におけるその２個頂点のindexの差を取り、その値を頂点毎に総和を取ったものが数列のスコアとなる。

ここで、1～Nのpermutationの形を成すN!通りの数列を考える。
このうち、上記計算方法でスコアがSとなるものの数を求めよ。

解法

まず計算量を無視して考える。
数列から生成される木について、dp[要素数][根が数列中何番目にあるか][部分木のスコア] := 条件を満たす木の数、というDPを考えよう。
dp[N][L][S]は以下のように計算できる

  • N==1の時、スコアは0にしかならないのでdp[N][L][S]はS=0なら1、S!=0なら0
  • L==1またはL==Nの時、この頂点は子を1個しか持たないので、残り(N-1)頂点の事だけ考えればよく、残り(N-1)頂点の根がどこに来るかを考慮して、dp[N][L][S] := sum_i(dp[N][i][S])
  • 2≦L≦N-1の時、この頂点は子を２個持つ。左側の子頂点が数列中x番目、右側の子頂点数列中L+y番目にあるとすると、その２子頂点によりスコアが(L+y)-x点上昇する。
    • よって、(N-1)頂点を左に(L-1)個、右に(N-L)個振り分けることも考えると、dp[N][L][S] = sum(Comb(N-1,L-1) * dp[L-1][x][a] * dp[N-L][y][b]) ただしa,b,x,yはa+b+(L+y)-x == Sを満たす範囲で総当たりとなる。

最後の計算が問題で、DPの状態がそもそもO(N^2*S)程度あるのに、最後３つの変数をループさせると全部でO(N^4*S^2)もかかりとても間に合わない。
そこで、最後の３変数ループの計算量を落とすことを考える。
上記式だとサブツリーの根の位置x,yと、サブツリーのスコアa,bを別々に総当たりしているが、位置x,yは結局最終的にスコア計算に用いるものである。
この考察より、位置情報を最初からスコアに盛り込みたい。

ここで左側のサブツリーについてもう少し考える。
dp[L-1][x][a]はサブツリー自体のスコアはaである。
ただ、このサブツリーが親頂点の左側に接続された場合、その時点で親ツリーのスコアに(L-x)だけ加算されるのが明らかである。
だとすると、最初から「このサブツリーは、親頂点の左側に接続すると(L-x)+a点スコアが加算されるサブツリーだ」と考えると、位置を考慮しなくてよくなる。

例えばdpl[N][S] := (親頂点の左側に接続するとS点スコアが加算されるようなN頂点のツリーの数)とする。
同様にdpr[N][S] := (親頂点の右側に接続するとS点スコアが加算されるようなN頂点のツリーの数)

こうするとさっきの式はdp[N][L][S] = sum(Comb(N-1,L-1) * dpl[L-1][a] * dp[N-L][b]) をa+b==Sを満たす範囲で計算すればよくなり、計算量はO(N^2*S^2)まで落とすことができる。
なお、dpl[N][S] := sum_i(dp[N][i][S-(N-i)]で計算できる。

以下のコードでは、左右の対称性からdpl=dprと見なし、dplだけ計算している。

ll memo[101][55][101];
ll lmemo[101][101];
const int CN=1001;
ll C[CN][CN];

class BearPermutations {
	public:
	ll mo;
	
	ll glmemo(int N,int S) {
		
		if(lmemo[N][S]>=0) return lmemo[N][S];
		
		ll ret=0;
		for(int i=0;i<N;i++) {
			int sc=S-(1+i);
			if(sc>=0) ret+=dpdp(N,sc,i);
		}
		return lmemo[N][S]=ret%mo;
		
	}
	
	ll dpdp(int N,int S,int left) {
		if(N<=1) return S==0;
		if(left>=N/2) left=N-1-left;
		if(memo[N][left][S]>=0) return memo[N][left][S];
		
		ll ret=0;
		if(left==0) {
			for(int i=1;i<N;i++) ret+=dpdp(N-1,S,i-1);
		}
		else {
			for(int ls=0;ls<=S;ls++) (ret+=glmemo(left,ls)*glmemo(N-left-1,S-ls))%=mo;
			(ret*=C[N-1][left])%=mo;
		}
		return memo[N][left][S]=ret%mo;
	}
	
	
	int countPermutations(int N, int S, int MOD) {
		mo = MOD;
		int i,j;
		
		FOR(i,CN) for(j=0;j<=i;j++) C[i][j]=(j==0||j==i)?1:(C[i-1][j-1]+C[i-1][j])%mo;
		
		MINUS(lmemo);
		MINUS(memo);
		ll ret=0;
		
		FOR(i,S+1) FOR(j,N) ret+=dpdp(N,i,j) % mo;
		return ret%mo;
	}
}

まとめ

式を眺めていたら計算量が落とせたので良かった。