[国家集训队]聪聪可可
题目链接:点这里
大致题意:给你一颗树,求出所有边权和能被3整除的路径数
Solution:
用一个pos数组来记录路径长度mod 3之后的数量
则答案为pos[1] * pos[2] * 2+pos[3]*pos[3]
最后去除同子树内的答案即可
Code:
#include<bits/stdc++.h>
#define ll long long
#define inf 2147483647
using namespace std;
const int N=2e4+1;
int rt,cnt,szt,head[N],sz[N],mx[N];
int n,ans,tot,dis[N],vis[N],pos[3];
struct Edge{int nxt,to,val;}edge[N<<1];
void ins(int x,int y,int z){
edge[++cnt].nxt=head[x];
edge[cnt].to=y;edge[cnt].val=z;
head[x]=cnt;
}
int gcd(int x,int y){return y==0?x:gcd(y,x%y);}
void getrt(int x,int f){
sz[x]=1,mx[x]=0;
for(int i=head[x];i;i=edge[i].nxt){
int y=edge[i].to;
if(y==f||vis[y]) continue;
getrt(y,x);sz[x]+=sz[y];
mx[x]=max(mx[x],sz[y]);
}mx[x]=max(mx[x],szt-sz[x]);
if(mx[x]<mx[rt]) rt=x;
}
void getdis(int x,int f,int d){
dis[++tot]=d;pos[d%3]++;
for(int i=head[x];i;i=edge[i].nxt){
int y=edge[i].to;
if(y==f||vis[y]) continue;
getdis(y,x,d+edge[i].val);
}
}
int calc(int x,int d){
pos[0]=pos[1]=pos[2]=0;
tot=0;getdis(x,0,d);
return pos[1]*pos[2]*2+pos[0]*pos[0];
}
void divide(int x){
ans+=calc(x,0);vis[x]=1;
for(int i=head[x];i;i=edge[i].nxt){
int y=edge[i].to;
if(!vis[y]){
ans-=calc(y,edge[i].val);
rt=0,szt=sz[y],getrt(y,0);
divide(rt);
}
}
}
int read(){
int x=0,f=1;char ch=getchar();
while(!isdigit(ch)){if(ch=='-')f=-f;ch=getchar();}
while(isdigit(ch)){x=x*10+ch-48;ch=getchar();}
return x*f;
}
int main(){
n=read();
for(int i=1;i<n;i++){
int x=read(),y=read(),z=read();
ins(x,y,z);ins(y,x,z);
}
szt=n;mx[rt]=inf;
getrt(1,0);divide(rt);
int d=gcd(ans,n*n);
printf("%d/%d",ans/d,n*n/d);
return 0;
}