2019沈阳网络赛 D.Fish eating fruit (点分治)

题意：

求树上任意点对之间路径长度%3=0,1,2的路径和

思路：

统计想歪了没弄出来，看了别人的题解说要统计 %3=0,1,2 的路径总长度和 %3=0,1,2 出现的次数
开始没懂为什么要统计次数，画个图就明白了：

d[i]为 %3=i 的路径条数
sum[i]为 %3=i 的路径总长度

该图中:
d[1]=2,d[2]=1;
sum[1]=2,sum[2]=2;
简称路径i为 %3=i 的路径
因为路径2要和d[1]条路径1匹配，所以路径2总贡献为d[1]*sum[2];
同理路径1的总贡献为d[2]*sum[1];

code：

//https://nanti.jisuanke.com/t/41403
#include<cstdio>
#include<cstring>
#include<vector>
#include<iostream>
#include<algorithm>
using namespace std;
#define ll long long
const int maxm=3e4+5;
const int mod=1e9+7;
struct mp{
    int v,w;
    mp(int vv,int ww){
        v=vv,w=ww;
    }
};
vector<mp>g[maxm];
int sz[maxm],son[maxm];
ll sum[3],d[3],ans[3];//sum[]是路径和,d[]是路径条数
int mark[maxm];
int n,size,root;
void init(){
    for(int i=1;i<=n;i++){
        g[i].clear();
        mark[i]=0;
    }
    ans[0]=ans[1]=ans[2]=0;
}
void getroot(int x,int fa){
    sz[x]=1;
    son[x]=0;
    int len=g[x].size();
    for(int i=0;i<len;i++){
        int v=g[x][i].v;
        if(mark[v]||v==fa)continue;
        getroot(v,x);
        sz[x]+=sz[v];
        son[x]=max(son[x],sz[v]);
    }
    son[x]=max(son[x],size-sz[x]);
    if(son[x]<son[root]){
        root=x;
    }
}
void getd(int x,int val,int fa){
    d[val%3]++;
    sum[val%3]+=val;
    sum[val%3]%=mod;
    int len=g[x].size();
    for(int i=0;i<len;i++){
        int v=g[x][i].v;
        int w=g[x][i].w;
        if(mark[v]||v==fa)continue;
        getd(v,(val+w)%mod,x);
    }
}
void solve(int x,int val,int change){
    sum[0]=sum[1]=sum[2]=0;
    d[0]=d[1]=d[2]=0;
    getd(x,val%mod,0);
    for(int i=0;i<3;i++){
        for(int j=0;j<3;j++){
            ans[(i+j)%3]=(ans[(i+j)%3]+d[i]*sum[j]*change%mod+mod)%mod;
            ans[(i+j)%3]=(ans[(i+j)%3]+d[j]*sum[i]*change%mod+mod)%mod;
        }
    }
}
void divide(int x){
    solve(x,0,1);
    mark[x]=1;
    int len=g[x].size();
    for(int i=0;i<len;i++){
        int v=g[x][i].v;
        int w=g[x][i].w;
        if(mark[v])continue;
        solve(v,w%mod,-1);
        son[root=0]=size=sz[v];
        getroot(v,0);
        divide(root);//这里老是习惯性错写成divide(v),导致tle
    }
}
signed main(){
    while(scanf("%d",&n)!=EOF){
        init();
        for(int i=1;i<n;i++){
            int a,b,c;
            scanf("%d%d%d",&a,&b,&c);
            a++,b++;
            g[a].push_back(mp(b,c));
            g[b].push_back(mp(a,c));
        }
        son[root=0]=size=n;
        getroot(1,0);
        divide(root);
        printf("%lld %lld %lld\n",ans[0],ans[1],ans[2]);
    }
    return 0;
}