AtCoder ABC #133 : E - Virus Tree 2

もう少しだったんだけどな。
https://atcoder.jp/contests/abc133/tasks/abc133_e

問題

N頂点の木を成す無向グラフが与えられる。
これらの頂点をK色で塗り分けたい。
その際、距離2以下に同じ色の頂点が来ないような塗り方は何通りか。

解法

DFS順で色を確定させていくことを考える。
その際、すでに処理が完了した距離2以内の頂点がi個ある場合、(K-i)を掛ける、ということを行っていこう。

DFSで処理済みの頂点のうち、距離が2以内の点は

• 親
• 親の親
• 兄(親の子のうち自身より探索は先に行われた頂点)

である。これらは親頂点がDFSを行う過程で数を子に伝えていけば実装が軽くなる。

int N,K;
vector<int> E[101010];
ll mo=1000000007;
ll ret=1;

void dfs(int cur,int pre,int num) {
	ret=ret*num%mo;
	
	num=K-1;
	if(pre!=-1) num--;
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur,num);
		num--;
	}
	
	
}

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

まとめ

Dより簡単な気がする。