TopCoder SRM 559 Div1 Medium HatRack

さて続いてMedium。
今回１ミスしたけど、ノーヒントでMediumが解けたのは良い感じだ。
http://community.topcoder.com/stat?c=problem_statement&pm=12176

木構造が与えられたとき、条件を満たすバランス木の数を求める。
木構造の各点を最上位の点としたときに、以下がバランス木になっているか、またバランス木のパターンを再帰的に求めていく。
各点の下部に２つサブツリーがあるとして、両者が同じ深さで同じバランス木だったら、どちらのサブツリーを前に持ってきても良いので、その点以下のパターンはサブツリーのパターンの積の倍になる。
両サブツリーの大きさに差があれば、大きい方を前に持ってくるしかないのでその点以下のパターンはサブツリーのパターン数になる。

class HatRack {
	public:
	int N;
	int MD,left;
	int ncon[52];
	int mat[52][52];
	
	ll npat[52][52];
	int nhat[52][52];
	int ndep[52][52];
	
	void dfs(int from, int cur) {
		int i,tar,t[2],nt;
		int first=(from==51)?1:0;
		if(ndep[from][cur]!=-1) return;
		
		if(ncon[cur]==1) {
			ndep[from][cur]=0;
			nhat[from][cur]=1;
			npat[from][cur]=1;
			goto out;
		}
		else if(ncon[cur]==2) {
			FOR(i,ncon[cur]) {
				tar = mat[cur][i];
				if(tar!=from) {
					dfs(cur,tar);
					if(ndep[cur][tar]!=0) {
						ndep[from][cur] = -2;
					}
					else {
						ndep[from][cur] = 1;
						npat[from][cur] = 1;
						nhat[from][cur] = 2;
					}
					goto out;
				}
			}
		}
		else if(ncon[cur]==3) {
			nt=0;
			FOR(i,ncon[cur]) {
				tar = mat[cur][i];
				if(tar == from) continue;
				t[nt++]=tar;
				dfs(cur,tar);
			}
			
			nhat[from][cur]=nhat[cur][t[0]]+nhat[cur][t[1]]+1;
			if(ndep[cur][t[0]]==-2 || ndep[cur][t[1]]==-2) {
				ndep[from][cur] = -2;
			}
			else if(ndep[cur][t[0]]==ndep[cur][t[1]]) {
				if(nhat[cur][t[0]] != (1<<(1+ndep[cur][t[0]]))-1 &&
				   nhat[cur][t[1]] != (1<<(1+ndep[cur][t[1]]))-1) {
					ndep[from][cur]=-2;
					goto out;
				}
				ndep[from][cur]=ndep[cur][t[0]]+1;
				if(nhat[cur][t[0]] == (1<<(1+ndep[cur][t[0]]))-1 &&
				   nhat[cur][t[1]] == (1<<(1+ndep[cur][t[1]]))-1) {
					npat[from][cur]=2*npat[cur][t[0]]*npat[cur][t[1]];
				}
				else {
					npat[from][cur]=npat[cur][t[0]]*npat[cur][t[1]];
				}
			}
			else if(ndep[cur][t[0]]==ndep[cur][t[1]]+1) {
				if(nhat[cur][t[1]] != (1<<(1+ndep[cur][t[1]]))-1) {
					ndep[from][cur]=-2;
				}
				else {
					ndep[from][cur]=ndep[cur][t[0]]+1;
					npat[from][cur]=npat[cur][t[0]]*npat[cur][t[1]];
				}
			}
			else if(ndep[cur][t[0]]+1==ndep[cur][t[1]]) {
				if(nhat[cur][t[0]] != (1<<(1+ndep[cur][t[0]]))-1) {
					ndep[from][cur]=-2;
				}
				else {
					ndep[from][cur]=ndep[cur][t[1]]+1;
					npat[from][cur]=npat[cur][t[0]]*npat[cur][t[1]];
				}
			}
			else {
				ndep[from][cur] = -2;
			}
		}
out:
		return;
	}
	
	
	long long countWays(vector <int> knob1, vector <int> knob2) {
		int i,j;
		N=knob1.size()+1;
		MD=0;
		while(N>=(1<<(MD+1))) MD++;
		left = N-((1<<(MD-1))-1);
		
		
		ZERO(ncon);
		ZERO(mat);
		FOR(i,knob1.size()) {
			int f = knob1[i]-1;
			int t = knob2[i]-1;
			mat[f][ncon[f]++] = t;
			mat[t][ncon[t]++] = f;
		}
		FOR(i,N) if(ncon[i]==0 || ncon[i]>=4) return 0;
		
		ZERO(npat);
		MINUS(nhat);
		MINUS(ndep);
		long long res=0;
		FOR(i,N) {
			mat[i][ncon[i]++]=51;
			dfs(51,i);
			ncon[i]--;
			if(ndep[51][i]==MD) res += npat[51][i];
		}
		
		return res;
		
		
	}

};

まとめ

せっかく解けたのは良かったが、ノーミスで行きたかったな。
今回はEasyのHyperKnightといい、計算量的には大したことない問題だね。