kmjp's blog

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

TopCoderOpen 2017 Round2B Medium DengklekGaneshAndTree

妙に実装に手間取った。
https://community.topcoder.com/stat?c=problem_statement&pm=14633

問題

根付き木が与えられる。
各頂点にはアルファベットが割り当てられている。
各頂点は大文字小文字のどちらかを割り当てたい。以下の条件を満たす割り当て方は何通りか。

  • 根からの深さが等しい頂点群に対し、最低1個は大文字の頂点がある。
  • 同じアルファベット同士を辺がつないでいる場合、両端の頂点は大文字小文字一致しなければならない。

解法

「最低1個は~~」の処理が面倒なので、全部小文字のケースを取り除いていこう。
重複を防ぐため、「深さがi未満の頂点には全部小文字の深さがなく、深さiの頂点群は全部小文字。深さ(i+1)以上は任意」というものを考えていく。

まず葉から順にDPしていき、以下を求めておこう。これはよくある木DPである。
D(v,c) := 頂点vが小文字であることが真偽値cである場合、後者の条件を満たすSubTree以下の頂点の大文字小文字の組み合わせ
深さがdである頂点群V(d)において、すべてのv∈V(d)に対しD(v,true)の積を取れば、深さdの頂点群がすべて小文字の場合の深さ(d+1)以降の頂点の組み合わせがわかる。
上記値をE(d) := v∈V(d)に対しD(v,true)の積 としておこう。

あとは、深さdまでに全部小文字の深さがないケースを求めればよい。すなわち以下を求めていく。
dp(d) := 深さdまでの部分木を考えたとき、深さdの頂点がすべて小文字で、深さd未満の頂点で同じ深さの頂点群がすべて小文字な深さがない組み合わせ。

これは各深さdに対し木DPを繰り返すとよい。
以下を考える。
F(v,d,c) := 深さdの頂点がすべて小文字の時、頂点vが小文字であることが真偽値cである場合、後者の条件を満たす深さdまでのSubTree以下の頂点の大文字小文字の組み合わせ

v∈V(d)に対しF(v,d,true)=1、F(v,d,false)=0として、深さd未満の頂点に対し同様に木DPしよう。
G(d',d) := 深さdの頂点がすべて小文字の時、深さd'の頂点がすべて小文字となるような深さd'~dの頂点の組み合わせ
上記G(d',d)はv∈V(d')に対しF(v,d,true)の積として求められる。
dp(d)は全ケース(F(root,d,true)+F(root,d,false))からd'で初めて全頂点小文字になるケースG(d',d)*dp(d')を順次引いていけばよい。

あとは全ケースD(root,false)+D(root,true)から、各深さdで初めて全頂点小文字になるケース(dp(d)*E(d))を引いていくだけ。

int N;
int D[1010];
vector<int> E[1010];

ll upper[1010];
ll up[1010][2];
ll down[1010][2];
ll cl[1010],dl[1010];

ll mo=1000000007;

class DengklekGaneshAndTree {
	public:
	int getCount(string labels, vector <int> parents) {
		
		int i;
		N=parents.size()+1;
		FOR(i,N) {
			E[i].clear();
			dl[i]=1;
		}
		FOR(i,N-1) {
			D[i+1]=D[parents[i]]+1;
			E[parents[i]].push_back(i+1);
		}
		
		for(i=N-1;i>=0;i--) {
			down[i][0]=down[i][1]=1;
			FORR(e,E[i]) {
				if(labels[i]==labels[e]) {
					(down[i][0]*=down[e][0])%=mo;
					(down[i][1]*=down[e][1])%=mo;
				}
				else {
					(down[i][0]*=down[e][0]+down[e][1])%=mo;
					(down[i][1]*=down[e][0]+down[e][1])%=mo;
				}
			}
			(dl[D[i]]*=down[i][0])%=mo;
		}
		
		ll ret=down[0][0]+down[0][1];
		ZERO(upper);
		
		int lv;
		FOR(lv,N) {
			ZERO(up);
			
			int ok=0;
			FOR(i,N) if(D[i]==lv) up[i][0]=ok=1;
			if(ok==0) break;
			
			FOR(i,N) cl[i]=1;
			for(i=N-1;i>=0;i--) if(D[i]<lv) {
				up[i][0]=up[i][1]=1;
				FORR(e,E[i]) {
					if(labels[i]==labels[e]) {
						(up[i][0]*=up[e][0])%=mo;
						(up[i][1]*=up[e][1])%=mo;
					}
					else {
						(up[i][0]*=up[e][0]+up[e][1])%=mo;
						(up[i][1]*=up[e][0]+up[e][1])%=mo;
					}
				}
				(cl[D[i]]*=up[i][0])%=mo;
			}
			
			upper[lv]=(up[0][0]+up[0][1])%mo;
			for(i=lv-1;i>=0;i--) (upper[lv]+=mo-cl[i]*upper[i]%mo)%=mo;
			(ret += mo-upper[lv]*dl[lv]%mo)%=mo;
		}
		
		return (ret%mo+mo)%mo;
		
	}
}

まとめ

上から下からDPして割としんどい。