CSAcademy Round #51 : D. Tree Coloring

Div2ということを踏まえても簡単だった気がする。
https://csacademy.com/contest/round-51/task/tree-coloring/

問題

N頂点の木を成す無向グラフの各頂点をK色のいずれかで塗る。
この時、距離が２以内の頂点同士が同じ色とならないようにしたい。
塗り分け方は何通りか。

解法

根付き木とみなし、根頂点から順に塗っていこう。
根はK通りとして、以後根以外の頂点を考える。

親から順に色を塗るので、子頂点側は考えなくてよいとするとき、各頂点は何通りの塗り方が可能か。
まず、親頂点および親の親とは同じ色を塗れないので、1ないし2色候補が減る。
そのほか、親を同じくする頂点同士も同じ色で塗れない。
よってそのような頂点同士の色を順に塗るとき、１つずつ選択肢が減っていく。

後は全頂点における選択肢の積を取るだけ。
DFSこそ行うものの複雑なDP等は不要である。

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

ll ret=1;

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

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,-1,K);
	cout<<ret<<endl;
	
}

まとめ

考察にちょっと手間取った。