SCOI2019 RGB [容斥]

传送门
考虑只有一个绿点的情况，就是一个裸的树形 dp，强制选当前点
$f_{u,r}=\prod (f_{v,r}+1),f_{u,b}=\prod (f_{v,b}+1)$
$ans=f_{u,r}\ast f_{u,b}$
我们可以枚举所有绿点算一遍答案，考虑有哪些情况会算重
如果绿点不相连，很明显不会重，会重的情况只有绿色点聚成一坨的时候
我们要让一坨的点的贡献只算一次，发现点数-边数=1
于是可以枚举每个绿点，在减去每条边的贡献就可以让一个连通块只算一次了

#include<bits/stdc++.h>
#define N 4050
using namespace std;
const int Mod = 1000000007;
int read(){
	int cnt = 0, f = 1; char ch = 0;
	while(!isdigit(ch)){ ch = getchar(); if(ch == '-') f = -1;}
	while(isdigit(ch)){ cnt = cnt * 10 + (ch-'0'), ch = getchar();}
	return cnt * f;
}
typedef long long ll;
ll add(ll a, ll b){ return (a + b) % Mod;}
ll mul(ll a, ll b){ return (a * b) % Mod;}
int first[N], nxt[N], to[N], w[N], tot;
void add(int x, int y, int z){
	nxt[++tot] = first[x], first[x] = tot, to[tot] = y, w[tot] = z;
}
int n, W, col[N];
struct Edge{ int u, v, w;} E[N];
ll ans, fR[N], fB[N];
void dfs(int u, int fa, int dis){
	if(dis > W){ fR[u] = fB[u] = 0; return;}
	if(col[u] != 1) fB[u] = 1;
	if(col[u] != 3) fR[u] = 1;
	for(int i = first[u]; i; i = nxt[i]){
		int t = to[i]; if(t == fa) continue;
		dfs(t, u, dis + w[i]);
		fB[u] = mul(fB[u], fB[t] + 1);
		fR[u] = mul(fR[u], fR[t] + 1);
	}
}
void Yolanda(int u){ dfs(u, 0, 0); ans = add(ans, mul(fR[u], fB[u]));}
void FSY(int u, int v, int w){
	dfs(u, v, w); dfs(v, u, w);
	ans = add(ans, Mod - mul(mul(fR[u], fB[u]), mul(fR[v], fB[v])));
}
int main(){
	n = read(), W = read(); string s; cin >> s;
	for(int i = 0; i < n; i++){
		if(s[i] == 'R') col[i+1] = 1;
		if(s[i] == 'G') col[i+1] = 2;
		if(s[i] == 'B') col[i+1] = 3; 
	}
	for(int i = 1; i < n; i++){
		int x = read(), y = read(), z = read();
		add(x, y, z); add(y, x, z);
		E[i] = (Edge){ x, y, z };
	} 
	for(int i = 1; i <= n; i++) if(col[i] == 2) Yolanda(i);
	for(int i = 1; i < n; i++) if(col[E[i].u] == 2 && col[E[i].v] == 2) FSY(E[i].u, E[i].v, E[i].w);
	cout << ans; return 0;
}