kmjp's blog

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

yukicoder : No.1075 木の上の山

誤読してタイムロスした。
https://yukicoder.me/problems/no/1075

問題

N頂点の木を成す無向グラフの各頂点に1~Kの値を振る。
この時、ある点vを1つ定めると、そこからどの点へのパスにおける頂点の値も、広義単調減少とできそうな振り方は何通りか。

解法

vを定めて、vから離れるほど単調減少になるケースを数え上げればいいのだが、複数同じ値の点がある場合、vの候補が複数あり得る場合がある。
その場合、最も親に近いものをvとするものとする。

dp(v,k) := vの頂点の値をkとし、vのsubtreeの頂点に対し、値が広義単調減少になるような値の振り方
dp2(v,k) := vの親頂点の値をk未満とし、v親方向の頂点群に対し、値が広義単調減少になるような 値の振り方
とすると、dp(v,k)*dp2(v,k)の総和が解。
dpは1回DFSを行って埋め、dp2はdpを用いて全方位DPの要領で親から求めて行こう。

int N,K;
vector<int> E[1010];
const ll mo=1000000007;
ll dp[1010][1010];
int C[1010];
ll ret;
ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}


int dfs(int cur,int pre) {
	int i;
	C[cur]=1;
	FOR(i,K) dp[cur][i]=1;
	FORR(e,E[cur]) if(e!=pre) {
		C[cur]+=dfs(e,cur);
		ll sum=0;
		FOR(i,K) {
			(sum+=dp[e][i])%=mo;
			(dp[cur][i]*=sum)%=mo;
		}
	}
	return C[cur];
}

void dfs2(int cur,int pre,vector<ll> P) {
	ll sum=0;
	int i;
	if(pre==-1) {
		FOR(i,K) ret+=dp[cur][i];
	}
	else {
		FOR(i,K) {
			(ret+=dp[cur][i]*sum)%=mo;
			(sum+=P[i])%=mo;
		}
		sum=0;
		FOR(i,K) {
			(sum+=P[i])%=mo;
			(dp[cur][i]*=sum)%=mo;
		}
	}
	
	FORR(e,E[cur]) if(e!=pre) {
		sum=0;
		FOR(i,K) {
			(sum+=dp[e][i])%=mo;
			P[i]=dp[cur][i]*modpow(sum)%mo;
		}
		dfs2(e,cur,P);
		
	}
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0,0);
	vector<ll> P(K);
	
	dfs2(0,-1,P);
	
	cout<<ret%mo<<endl;
}

まとめ

候補に迷うときは根に近い方に決める、というテクは応用ききそう。