「SCOI 2019 D2T2」RGB

传送门

---

problem

有一棵树，$n$ 个节点，每条边有边权 $c_i$，每个点的颜色为 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)\le w$。

答案对 $10^9+7$ 取模。

数据范围：$n\le 2000$，$w\le 10^7$，$c_i\le 10^6$。

---

solution

听说这是一道论文题。

我们先考虑只有一个绿点的时候怎么做。

考虑把绿点当作根，然后树形 $dp$。设 $\text{R}(x)$ 表示强制选 $x$ 时，以 $x$ 为根的子树中只有红点和绿点的连通块个数。$\mathtt{B}(x)$ 的定义与之类似：

$$\begin{array}{rl}\mathtt{R}(x)& =\prod _{v\in son(x)}(\text{R}(v)+1)\\ \mathtt{B}(x)& =\prod _{v\in son(x)}(\text{B}(v)+1)\end{array}$$

意思就是，$x$ 的儿子 $v$ 都有 $\text{R}(v)+1$ 种选择（$+1$ 是因为可以不选），由于强制选 $x$，用乘法原理乘一下就是答案。

那么答案就是 $\text{R}(r)\times \text{B}(r)$，$r$ 就是那个作为根的绿点。

那么如果有几个绿点是连通的，答案会被算重，考虑如果减去重复的部分。

这时想到，对于连通的一些绿点，点数减去边数等于 $1$。因此我们可以先枚举点加上，再枚举边减去。

时间复杂度 $O(n^2)$。

---

code

#include<bits/stdc++.h>
using namespace std;
const int N=4005,P=1e9+7;
int n,lim,t,ans,col[N];
int first[N],v[N],w[N],nxt[N];
struct edges{int u,v,w;}e[N];
int add(int x,int y)  {return x+y>=P?x+y-P:x+y;}
int dec(int x,int y)  {return x-y< 0?x-y+P:x-y;}
int mul(int x,int y)  {return 1ll*x*y%P;}
void edge(int x,int y,int z){
	nxt[++t]=first[x],first[x]=t,v[t]=y,w[t]=z;
}
int R[N],B[N];
void dfs(int x,int fa,int dis){
	if(dis>lim) {R[x]=B[x]=0;return;}
	R[x]=(col[x]!=3),B[x]=(col[x]!=1);
	for(int i=first[x];i;i=nxt[i]){
		int to=v[i];
		if(to==fa)  continue;
		dfs(to,x,dis+w[i]);
		R[x]=mul(R[x],R[to]+1);
		B[x]=mul(B[x],B[to]+1);
	}
}
void calc(int x){
	dfs(x,0,0);
	ans=add(ans,mul(R[x],B[x]));
}
void calc(int u,int v,int w){
	dfs(u,v,w),dfs(v,u,w);
	ans=dec(ans,mul(mul(R[u],B[u]),mul(R[v],B[v])));
}
char S[N];
int main(){
	scanf("%d%d%s",&n,&lim,S+1);
	for(int i=1;i<=n;++i){
		if(S[i]=='R')  col[i]=1;
		if(S[i]=='G')  col[i]=2;
		if(S[i]=='B')  col[i]=3;
	}
	for(int i=1,x,y,z;i<n;++i){
		scanf("%d%d%d",&x,&y,&z);
		edge(x,y,z),edge(y,x,z),e[i]=(edges){x,y,z};
	}
	for(int i=1;i<=n;++i)  if(col[i]==2)  calc(i);
	for(int i=1;i<n;++i){
		int u=e[i].u,v=e[i].v,w=e[i].w;
		if(col[u]==2&&col[v]==2)  calc(u,v,w);
	}
	printf("%d\n",ans);
	return 0;
}