Fish eating fruit 点分治

传送门：https://nanti.jisuanke.com/t/41403

求树中两点距离%3为0,1,2的距离和

#include
#include
#include
#include
using namespace std;
typedef long long ll;

#define N 10010
#define inf 1e9+10

const ll mod=1e9+7;

struct node{
	int to,c,next;
}g[N*2]; 

int head[N],m;
int size[N],f[N];//f:最大子树节点数
bool vis[N];
int d[N],deep[N];
int ct[4];
ll all[4],endall[4];
ll ans,ans1,ans2;
int n,sum,root,k;

void add_edge(int from,int to,int cost){
	g[++m].next=head[from];
	head[from]=m;
	g[m].to=to;g[m].c=cost;
}

void getroot(int v,int fa){
	size[v]=1;f[v]=0;
	for(int i=head[v];i;i=g[i].next){
		if(g[i].to!=fa&&!vis[g[i].to]){
			getroot(g[i].to,v);
			size[v]+=size[g[i].to];
			f[v]=max(f[v],size[g[i].to]);
		}
	}
	f[v]=max(f[v],sum-size[v]);
	if(f[v]<f[root]) root=v;
}

void getdeep(int v,int fa){
	deep[++deep[0]]=d[v];
	ct[d[v]%3]++;
	all[d[v]%3]+=d[v];	
	for(int i=head[v];i;i=g[i].next){
		if(g[i].to!=fa&&!vis[g[i].to]){
			d[g[i].to]=d[v]+g[i].c;
			getdeep(g[i].to,v);
		}
	}
}


void cal(int v,int cost){
	ct[0]=ct[1]=ct[2]=0;
	all[0]=all[1]=all[2]=0;
	d[v]=cost;
	deep[0]=0;
	getdeep(v,0);
	for(int i=0;i<=2;i++){
		for(int j=0;j<=2;j++){
			endall[(i+j)%3]+=ct[i]*all[j]+ct[j]*all[i];
		}
	}
}

void solve(int v){
	vis[v]=1;
	endall[0]=endall[1]=endall[2]=0;	
	cal(v,0);
	ans+=endall[0];
	ans1+=endall[1];
	ans2+=endall[2];
	ans%=mod;
	ans1%=mod;
	ans2%=mod;
	for(int i=head[v];i;i=g[i].next){
		if(!vis[g[i].to]){
			endall[0]=endall[1]=endall[2]=0;			
			cal(g[i].to,g[i].c);
			ans-=endall[0];
			ans1-=endall[1];
			ans2-=endall[2];
			ans=(ans+mod)%mod;
			ans1=(ans1+mod)%mod;
			ans2=(ans2+mod)%mod;			
			sum=size[g[i].to];
			root=0;
			getroot(g[i].to,0);
			solve(root);
		}
	}
}

int main(){
	int u,v,w;
	while(~scanf("%d",&n)&&n){
		ans=ans2=ans1=root=m=0;
		memset(vis,0,sizeof(vis));
		memset(head,0,sizeof(head));
		for(int i=1;i<n;i++){
			scanf("%d%d%d",&u,&v,&w);
			u++;
			v++;
			add_edge(u,v,w);
			add_edge(v,u,w);
		}
		f[0]=inf;
		sum=n;
		getroot(1,0);
		solve(root);
		printf("%lld %lld %lld\n",ans%mod,ans1%mod,ans2%mod);
	}
}

#include
#include
#include
#include
using namespace std;

#define rep(a,b,c) for(int a=b;a<=c;a++)
void _swap(int &x,int &y){x^=y;y^=x;x^=y;}
int _max(int x,int y){return (x>y)?x:y;}
inline int read(){
    int re=0,flag=1;char ch=getchar();
    while(ch>'9'||ch<'0'){
        if(ch=='-') flag=-1;
        ch=getchar();
    }
    while(ch>='0'&&ch<='9') re=(re<<1)+(re<<3)+ch-'0',ch=getchar();
    return re*flag;
}
const long long mod=1e9+7;
int n,cnt,ans,dis[200010],first[20010],tmp[5],root,siz[200010],son[200010],sum;
long long tempans[4],ans0,ans2,ans1,mo[4];
bool vis[20010];
struct edge{
    int to,next,w;
}a[40010];
inline void add(int u,int v,int w){
    a[++cnt]=(edge){v,first[u],w};first[u]=cnt;
    a[++cnt]=(edge){u,first[v],w};first[v]=cnt;
}
void init()
{
    cnt=0;memset(first,-1,sizeof(first));
}
int gcd(int x,int y){return (y?gcd(y,x%y):x);};
void getroot(int u,int f){
    int i,v;
    siz[u]=1;son[u]=0;
    for(i=first[u];~i;i=a[i].next){
        v=a[i].to;
        if(v==f||vis[v]) continue;
        getroot(v,u);
        siz[u]+=siz[v];
        son[u]=_max(son[u],siz[v]);
    }
    son[u]=_max(son[u],sum-siz[u]);
    if(son[u]<son[root]) root=u;
}

void gettmp(int u,int f){
    int i,v;
    tmp[dis[u]%3]++;
    mo[ dis[u]%3 ]+=dis[u];

    for(i=first[u];~i;i=a[i].next){
        v=a[i].to;
        if(v==f||vis[v]) continue;

        dis[v]=(dis[u]+a[i].w);
        gettmp(v,u);
    }
}
void calc(int u,int d){
    mo[0]=mo[1]=mo[2]=0;
    tmp[0]=tmp[1]=tmp[2]=0;
    
    dis[u]=d;
    gettmp(u,0);
    rep(i,0,2)
    {
        rep(j,0,2)
        tempans[ (i+j)%3 ] += mo[i] * tmp[j] + mo[j] * tmp[i];
    }
}
void dfs(int u){
    int i,v;
    vis[u]=1;
    tempans[0]=tempans[1]=tempans[2]=0;
    calc(u,0);
    ans1+=tempans[1];
    ans0+=tempans[0];
    ans2+=tempans[2];
    ans0%=mod;
    ans1%=mod;
    ans2%=mod;
    for(i=first[u];~i;i=a[i].next){
        v=a[i].to;
        if(vis[v]) continue;
        tempans[0]=tempans[1]=tempans[2]=0;
        calc(v,a[i].w);
        ans1-=tempans[1];
        ans0-=tempans[0];
        ans2-=tempans[2]; 
        ans0=(ans0+mod)%mod;
        ans1=(ans1+mod)%mod;
        ans2=(ans2+mod)%mod;
        sum=siz[v];root=0;
        getroot(v,u);
        dfs(root);
    }
}
int main(){

    while(scanf("%d",&n)==1)
    {
        init();ans1=ans2=ans0=0;
        memset(vis,0,sizeof vis);
        memset(first,-1,sizeof(first));
        int i,t1,t2,t3;
        for(i=1;i<n;i++){
            t1=read();t2=read();t3=read();t1++;t2++;
            add(t1,t2,t3);
        }
        sum=n;son[0]=n;
        getroot(1,0);
        dfs(1);
        printf("%lld %lld %lld\n",ans0%mod,ans1%mod,ans2%mod);
    }
}