CF1156D 题解

教练把这题出在了模拟赛。

很快啊，上来一眼感觉是点分治板子题，就赶紧写，写完了之后一直 WA on test #4，调了三个小时没调出来，气人，后来同学用各种方法切了，我还在沉迷于淀粉质无法自拔。

以上为闲话。

看到题目描述中的“路径”，很容易想到点分治。点分治的板子我就不说了，重点说说统计路径条数的部分。

假设我们当前遍历到的路径是以 $u$ 为根的子树，我们统计下面几种东西：

1. 从根到当前子树的点，边全部是 $0$ 的路径条数，记为 $c0$。
2. 从根到当前子树的点，边全部是 $1$ 的路径条数，记为 $c1$。
3. 从根到当前子树的点，边权先 $0$ 后 $1$ 的路径条数，记为 $c01$。
4. 从根到当前子树的点，边权先 $1$ 后 $0$ 的路径条数，记为 $c10$。

对于已经遍历过的子树，同样地统计出 $t0,t1,t01,t10$。

考虑如何统计答案，分类讨论。

首先，根节点也要统计在内，所以答案要加上 $2\times c0$（从$c0$走到根 和 从根走到 $c0$）。

同理，答案还要加上 $2\times c1,c01,c10$。

1. 路径的起点在 $c0$，合法的路径终点有 $t0$、$t1$、$t01$，所以答案要加上 $c0\times (t0+t1+t01)$。
2. 路径的起点在 $c1$，合法的路径终点只有 $t1$，所以答案要加上 $c1\times t1$。
3. 路径的起点在 $c10$，合法的路径终点只有 $t1$，所以答案要加上 $c10\times t1$。
4. 路径的起点在 $c01$，没有合法的路径终点（我就挂在这）。

对于路径的起点在 $t0,t1,t10,t01$ 的情况，同样地，答案要加上 $t0\times (c0+c1+c01)+t1\times c1+t10 \times c1$。

这样我们就 优雅地 解决了这道题。

最后强调一下不要求错重心。

Code

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N=2e5+10,INF=0x3f3f3f3f3f3f3f3f;
int n;
struct edge{
	int v,w,nxt;
}e[N*2];
int head[N],cnt=1;
void add(int u,int v,int w){
	e[cnt].v=v;
	e[cnt].w=w;
	e[cnt].nxt=head[u];
	head[u]=cnt++;
} 
int ans=0;
bool vis[N];
int root,mnroot;
int siz[N],s;
void getroot(int u,int fa){
	siz[u]=1;
	int mx=0;
	for(int i=head[u];i;i=e[i].nxt){
		if(e[i].v!=fa&&!vis[e[i].v]){
			getroot(e[i].v,u);
			siz[u]+=siz[e[i].v];
			mx=max(mx,siz[e[i].v]);
		}
	}
	mx=max(mx,s-siz[u]);
	if(mx<mnroot){
		mnroot=mx;
		root=u;
	}
}
int c0,c01,c10,c1;
void getdis(int u,int s1,int s2,int fa){
	if(s2==-1){
		if(s1)c1++;
		else c0++;
	}else{
		if(s1)c10++;
		else c01++;
	}
	for(int i=head[u];i;i=e[i].nxt){
		if(e[i].v!=fa&&!vis[e[i].v]){
			int t;
			if(e[i].w!=s1){
				if(s2==-1)t=e[i].w;
				else t=s2;
			}else{
				if(s2==-1)t=-1;
				else continue;
			}
			getdis(e[i].v,s1,t,u);
		}
	}
}
void calc(int u,int fa){
	int t0=0,t1=0,t01=0,t10=0;
	for(int i=head[u];i;i=e[i].nxt){
		if(!vis[e[i].v]&&e[i].v!=fa){
			c0=c1=c01=c10=0;
			getdis(e[i].v,e[i].w,-1,u);
			ans+=2*c0+2*c1+c01+c10;
			ans+=c0*(t1+t0+t01);
			ans+=c1*t1;
			ans+=c10*t1;
			ans+=t0*(c0+c1+c01);
			ans+=t1*c1;
			ans+=t10*c1;
			t1+=c1,t0+=c0,t01+=c01,t10+=c10;
		}
	}
}
void solve(int u,int fa){
	vis[u]=1;
	calc(u,fa);
	for(int i=head[root];i;i=e[i].nxt){
		if(!vis[e[i].v]&&e[i].v!=fa){
			s=siz[e[i].v];
			mnroot=INF;
			getroot(e[i].v,u);
			solve(root,u);
		}
	}
}
signed main(){
	ios::sync_with_stdio(false);
	cin.tie(0),cout.tie(0);
	cin>>n;
	for(int i=1;i<n;i++){
		int u,v,w;cin>>u>>v>>w;
		add(u,v,w);add(v,u,w);
	}
	s=n;
	mnroot=INF;
	getroot(1,0);
	solve(root,0);
	cout<<ans<<endl;
	return 0;
}