【树上计数 乘法原理】ABC129 E - Virus Tree 2

明天就要篮球被了，这不会是我今年最后一场比赛了吧 

E - Sum Equals Xor (atcoder.jp)

题意：

给定一棵树和K种颜色，让你给这棵树涂色，使得两点之间距离<=2的两个点不同色

思路： 

一开始很自然地想到了树形计数DP，但是很不幸，N,K<=1e5，爆空间了，因此考虑组合计数

第二个想法是容斥，先算出所有涂色方案是K^N，然后减去存在d=1的同色的，再减去d=2的同色的，然后加上d=1且d=2的同色的，然后发现这好像更难算了

然后就不会了，事实上，模拟一下第二个样例就会发现，这个就是乘法原理而已

一开始一定要先去模拟样例再去想思路！这样会省下很多时间

手摸一下第二个样例：

可以发现，深度为1的方案数是K，深度为2的方案数是K-1，其他的深度的最左边的点都是K-2，然后方案数随着从左到右递减

所以可以直接乘法原理了

Code：

#include <bits/stdc++.h>

#define int long long

using namespace std;

const int mxn=1e5+10;
const int mxe=1e5+10;
const int mod=1e9+7;

int N,K,u,v,tot=0;
int head[mxn];
int dep[mxn],a[mxn];

struct ty{
	int to,next;
}edge[mxe<<2];

void add(int u,int v){
	edge[tot].to=v;
	edge[tot].next=head[u];
	head[u]=tot++;
}
void G_init(){
	tot=0;
	for(int i=0;i<=N;i++){
		head[i]=-1;
	}
}
void dfs1(int u,int fa){
	dep[u]=dep[fa]+1;
	if(dep[u]==1) a[u]=K;
	else if(dep[u]==2) a[u]=K-1;
	else a[u]=K-2;
	for(int i=head[u];~i;i=edge[i].next){
		if(edge[i].to==fa) continue;
		dfs1(edge[i].to,u);
	}
}
void dfs2(int u,int fa){
	int ok=0,last;
	for(int i=head[u];~i;i=edge[i].next){
		if(edge[i].to==fa) continue;
		if(!ok){
			ok=1;
			last=a[edge[i].to];
		}else{
			a[edge[i].to]=last-1;
			last=a[edge[i].to];
		}
	}
	for(int i=head[u];~i;i=edge[i].next){
		if(edge[i].to==fa) continue;
		dfs2(edge[i].to,u);
	}
}
void solve(){
	cin>>N>>K;
	G_init();
	for(int i=1;i<=N-1;i++){
		cin>>u>>v;
		add(u,v);
		add(v,u);
	}
	dfs1(1,0);
	dfs2(1,0);
	//for(int i=1;i<=N;i++) cout<<a[i]<<" \n"[i==N];
	int ans=1ll;
	for(int i=1;i<=N;i++){
		int u=max(0ll,a[i]);
		ans=(ans*u)%mod;
	}
	cout<<ans<<'\n';
}
signed main(){
	ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
	int __=1;//cin>>__;
	while(__--)solve();return 0;
}