[国家集训队]聪聪可可

大意:求树上路径和mod 3 =0 的点对数。
一开始写的容斥无限智障,后来参考了一下题解,发现根本不用我那种垃圾写法。。。

#include <iostream>
#include <cstdio>
#include <cstring>
using namespace std;
const int N=100005;
int n,rt,head[N],ecnt,siz[N],f[N],size,cnt[5],sum[5],ans;struct Edge{int to,nxt,val;}e[N<<1];
bool vis[N];
void add(int bg,int ed,int val) {e[++ecnt].nxt=head[bg];e[ecnt].to=ed;e[ecnt].val=val;head[bg]=ecnt;}
int gcd(int x,int y) {return !y?x:gcd(y,x%y);}
void barycenter(int x,int fa) {
	f[x]=0,siz[x]=1;
	for(int i=head[x];i;i=e[i].nxt) {
		int v=e[i].to;
		if(vis[v]||v==fa) continue;
		barycenter(v,x);
		f[x]=max(f[x],siz[v]);
		siz[x]+=siz[v];
	}
	f[x]=max(f[x],size-siz[x]);
	if(f[rt]>f[x]) rt=x;	
}
void dfs1(int x,int fa,int val) {
	sum[val]++;
	for(int i=head[x];i;i=e[i].nxt) {
		int v=e[i].to;
		if(vis[v]||v==fa) continue;
		dfs1(v,x,(e[i].val+val)%3);
	}
}
void dfs2(int x,int fa,int val) {
	cnt[val]++;
	for(int i=head[x];i;i=e[i].nxt) {
		int v=e[i].to;
		if(vis[v]||v==fa) continue;
		dfs2(v,x,(e[i].val+val)%3);
	}
}
int work(int x,int val) {
	memset(sum,0,sizeof sum);
	dfs1(x,0,val);
	return sum[1]*sum[2]*2+sum[0]*sum[0];
}
void solve(int x) {
	ans+=work(x,0);vis[x]=1;
	for(int i=head[x];i;i=e[i].nxt) {
		int v=e[i].to;
		if(vis[v]) continue;
		ans-=work(v,e[i].val);
		rt=0,siz[rt]=size=siz[v],barycenter(v,0),solve(rt);
	}
}
int main() {
	scanf("%d",&n);
	for(int i=1,x,y,z;i<n;i++) {scanf("%d%d%d",&x,&y,&z);add(x,y,z%3);add(y,x,z%3);}
	f[rt]=size=n;
	barycenter(1,0);
	solve(rt);int g=gcd(ans,n*n);
	printf("%d/%d",((ans)/g),(n*n/g));
}
我是咸鱼。转载博客请征得博主同意Orz
原文地址:https://www.cnblogs.com/sdfzhsz/p/9476266.html