kmjp's blog

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

yukicoder : No.364 門松木

こちらは発想は大変だけど実装は少し軽め。
http://yukicoder.me/problems/1009

問題

各頂点vに値A[v]を持つ木が与えられる。
木が門松木であるとは、次数2以上の各頂点vに対し異なる2頂点u,wを取るとA[u], A[v], A[w]が門松列を持つものを言う。

入力の木に対する単一の部分木のうち、門松木を成すもので内包する門松列の数(u,v,wの組み合わせ)の最大値を求めよ。

解法

木DPで、DFSで子から順に各頂点を根とする部分木における門松列数の最大値を求めていく。
各頂点vを根とした部分木を作る場合、以下の条件を満たす範囲で子頂点集合を部分木に含めることができる。

  • 子頂点集合の値は、全てA[v]未満かすべてA[v]より大きいかいずれか。
  • 子頂点集合の値は、重複するものを含めてはならない。
    • 言い換えると、同じ値を持つ子頂点が複数ある場合、サブツリー内の門松列数が最大のものを選ぶのが良い。


各頂点vに対し、以下の値を更新していこう。
D[v][c] := vの子頂点wのうち値A[w]がcのもののうち、wのサブツリー内の門松列数の最大値
N_UP[v] := vの子頂点wのうち、値A[w]がA[v]を超える値の数(同じ値を持つ複数の頂点は1つと数える。)
N_DOWN[v] := 同A[v]未満の値の数
S_UP[v] := vの子頂点wのうち、値A[w]がA[v]を超える頂点を選んだ時のサブツリー内の門松列数の総和(同じA[w]を成す子頂点が複数ある場合、サブツリー内の門松列数が最大のものを選んだ場合)
S_DOWN[v] := 同A[v]未満の頂点を選んだ場合。

各頂点におけるS_UP[v]、S_DOWN[v]の最大値が解。
例えばS_UP[v]の求め方を考える。
S_UP[v]に含まれるのは、以下の総和である。

  • 今見ている頂点vを中心とする門松列の数。すなわちN_DOWN[v]個の隣接頂点のうち、2個を選ぶ選びかた。
  • 子頂点wのうちA[w]がA[v]を超えるものにおけるS_DOWN[w]の総和
    • ただし、S_DOWN[w]には隣接頂点にA[v]と同じものがすでに含まれている場合、A[v]をwに連結することができない。よってwの子頂点のうちA[v]と値が一致するものを切り離す。
      • よってS_DOWN[w]からD[w][A[r]](すなわち値がA[r]となる孫頂点以下のサブツリーの門松列数)だけ引いたものを考える。
    • A[w]が一致するものが複数ある場合、S_DOWN[w](もしくはS_DOWN[w]-D[w][A[r]])が最大となるものを選ぶ。

ややこしく見えるけどコードは短い。

int N;
int A[50505];
vector<int> E[50505];
int dp[50505][1010];
int num_down[50505], num_up[50505];
int sc_down[50505], sc_up[50505];
int ma;

int dfs(int cur,int pre) {
	int i, ret=0;
	
	FORR(r,E[cur]) if(r!=pre) {
		ret = max(ret,dfs(r,cur));
		if(A[cur]<A[r]) {
			if(dp[r][A[cur]]>=0) dp[cur][A[r]] = max(dp[cur][A[r]],sc_down[r]-dp[r][A[cur]]);
			else dp[cur][A[r]] = max(dp[cur][A[r]],sc_down[r]+num_down[r]);
		}
		if(A[cur]>A[r]) {
			if(dp[r][A[cur]]>=0) dp[cur][A[r]] = max(dp[cur][A[r]],sc_up[r]-dp[r][A[cur]]);
			else dp[cur][A[r]] = max(dp[cur][A[r]],sc_up[r]+num_up[r]);
		}
	}
	
	for(i=1;i<=1000;i++) if(dp[cur][i]>=0) {
		if(i<A[cur]) num_down[cur]++, sc_down[cur]+=dp[cur][i];
		if(i>A[cur]) num_up[cur]++, sc_up[cur]+=dp[cur][i];
	}
	
	sc_down[cur] += num_down[cur]*(num_down[cur]-1)/2;
	sc_up[cur] += num_up[cur]*(num_up[cur]-1)/2;
	return max(ret, max(sc_down[cur],sc_up[cur]));
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	MINUS(dp);
	
	cout<<dfs(0,-1)<<endl;
	
}

まとめ

今後まだ門松列問題は出るのだろうか…。