Fish eating fruit 点分治

树上任意两点之间的路径按照模 3 为 012 分类，将两点间距离加和（边权累和）  ，乘 2 即为答案。 

点分治模板题

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
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);
    }
}

View Code