ABC133 E Virus Tree 2 树形遍历+染色

题目：E - Virus Tree 2

给定一个含$K$个节点的无向树，给每个顶点赋上$[1,K]$一整数（染上一种颜色）。并且，若两个顶点距离$\le 2$，那么两顶点值不同（颜色不同）。

$1\le N,K\le 10^5$

考察

首先无向树的根可以随意设定，我们就从第1个开始查找（邻接表）。根的颜色肯定是K种可用。我们固定好根的颜色后，用DFS向下探索遍历所有的枝及叶。赋值的时候，我们可以考虑以下两点

• 根的颜色数为$K$
• 固定顶点值和父级节点值，其所有子节点的赋值方案为${}_{k-2}{P}_s$，其中$s$为子节点数量

代码实现

// some macroes and includes
//---------------------
#define MAXN 100005
#define MOD 1000000007
//---------------------

ll n,k;
vector<int> G[MAXN];
ll ret=1;

ll modmult(ll a , ll b ){
  return ( (a%MOD) * (b%MOD)  ) %MOD;
  }

void solve(ll lastId,ll curId,ll depth){
  	
	if(ret==0) {ret=0;return;}
	ll temp=depth==0 ? k:1;
	ll id=k-min(depth+1,ll(2));
	REP(i,G[curId].size()){
		if(G[curId][i] == lastId) continue;
		solve( curId , G[curId][i] , depth+1);
		temp = modmult(temp,max(id--,ll(0)));
	}
	ret = modmult(ret,temp);
	//cout<<ret<<endl;
}

int main(){
  	cin>>n>>k;
  	REP(i,n-1) {int a,b;cin>>a>>b;G[a].PB(b);G[b].PB(a);}
  	solve(-1,1,0);
	cout << ret <<endl;
	return 0;
}