AtCoder Beginner Contest 133 E - Virus Tree 2(组合数学)

题意

n个点的树k种颜色，距离不超过2的点对需颜色不同，求方案数

Code(copy)

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>

typedef long long LL;

const int N=100005;
const int MOD=1000000007;

int n,k,jc[N],ny[N],ans,cnt,last[N];
struct edge{int to,next;}e[N*2];

int A(int n,int m)
{
	return n<m?0:(LL)jc[n]*ny[n-m]%MOD;
}

void addedge(int u,int v)
{
	e[++cnt].to=v;e[cnt].next=last[u];last[u]=cnt;
	e[++cnt].to=u;e[cnt].next=last[v];last[v]=cnt;
}

void dfs(int x,int fa)
{
	int deg=0;
	for (int i=last[x];i;i=e[i].next)
	{
		if (e[i].to==fa) continue;
		dfs(e[i].to,x);
		deg++;
	}
	if (!fa) ans=(LL)ans*A(k,deg+1)%MOD;
	else ans=(LL)ans*A(k-2,deg)%MOD;
}

int main()
{
	scanf("%d%d",&n,&k);
	for (int i=1;i<n;i++)
	{
		int x,y;scanf("%d%d",&x,&y);
		addedge(x,y);
	}
	jc[0]=jc[1]=ny[0]=ny[1]=1;
	for (int i=2;i<=std::max(n,k);i++) jc[i]=(LL)jc[i-1]*i%MOD,ny[i]=(LL)(MOD-MOD/i)*ny[MOD%i]%MOD;
	for (int i=2;i<=std::max(n,k);i++) ny[i]=(LL)ny[i-1]*ny[i]%MOD;
	ans=1;
	dfs(1,0);
	printf("%d\n",ans);
	return 0;
}