NOIp训练 [SCOI2019]RGB（容斥原理）

传送门
题意：
有一棵树，n个节点，每条边有边权ci，每个点的颜色为R,G,B三种颜色之一（R为红色，G为绿色，B为蓝色）
你要统计有序对(U,V)的数量，其中U,V是两个点集
它还需要满足以下条件：
1、U和V都必须是连通的
2、U的颜色只能是红色或者绿色，V的颜色只能是绿色或者蓝色
3、存在一个U,V均包含的点x，使得对于所有在U,V集合中的点y，有dis(x,y)≤w
答案对10^9+7取模
思路：
最终选出来的绿色的点一定构成一个连通块。
显然如果枚举每一种绿色的连通块然后计算可以产生的红/蓝连通块个数是可以不重不漏的。
现在对于每一个绿点的连通块，我们想让它只被计算一次，由于这个绿点的连通块的点数-边数恰好等于1。
于是分别枚举绿点和连着两个绿点的边统计答案即可。
代码：

#include<bits/stdc++.h>
#define ri register int
using namespace std;
inline int read(){
	#define gc getchar
	int ans=0;
	char ch=gc();
	while(!isdigit(ch))ch=gc();
	while(isdigit(ch))ans=((ans<<2)+ans<<1)+(ch^48),ch=gc();
	return ans;
	#undef gc
}
typedef long long ll;
const int mod=1e9+7,N=2005;
inline int add(const int&a,const int&b){return a+b>=mod?a+b-mod:a+b;}
inline int dec(const int&a,const int&b){return a>=b?a-b:a-b+mod;}
inline int mul(const int&a,const int&b){return (ll)a*b%mod;}
inline void Add(int&a,const int&b){a=a+b>=mod?a+b-mod:a+b;}
inline void Dec(int&a,const int&b){a=a>=b?a-b:a-b+mod;}
inline void Mul(int&a,const int&b){a=(ll)a*b%mod;}
typedef pair<int,int> pii;
#define fi first
#define se second
vector<pii>e[N],S;
vector<int>V;
int n,lim,col[N],f[N][2],ans=0;
char s[N];
inline void dfs(int p,int fa,int dist){
	if(dist>lim){f[p][0]=f[p][1]=0;return;}
	f[p][0]=col[p]!=2;
	f[p][1]=col[p]!=1;
	for(ri i=0,v;i<e[p].size();++i){
		if((v=e[p][i].fi)==fa)continue;
		dfs(v,p,dist+e[p][i].se);
		Mul(f[p][0],add(f[v][0],1));
		Mul(f[p][1],add(f[v][1],1));
	}
}
inline void solve1(int p){dfs(p,0,0),Add(ans,mul(f[p][0],f[p][1]));}
inline void solve2(int p1,int p2,int W){
	dfs(p1,p2,W),dfs(p2,p1,W);
	int w1=mul(f[p1][0],f[p2][0]),w2=mul(f[p1][1],f[p2][1]);
	Dec(ans,mul(w1,w2));
}
int main(){
	n=read(),lim=read();
	scanf("%s",s+1);
	for(ri i=1;i<=n;++i)col[i]=s[i]=='R'?1:(s[i]=='G'?3:2);
	for(ri i=1,u,v,w;i<n;++i){
		u=read(),v=read(),w=read(),e[u].push_back(pii(v,w)),e[v].push_back(pii(u,w));
		if(col[u]+col[v]==6)S.push_back(pii(u,v)),V.push_back(w);
	}
	for(ri i=1;i<=n;++i)if(col[i]==3)solve1(i);
	for(ri i=S.size()-1;~i;--i)solve2(S[i].fi,S[i].se,V[i]);
	cout<<ans;
	return 0;
}