The Preliminary Contest for ICPC Asia Shenyang 2019 D. Fish eating fruit

题意：设f(x,y)=dis(x,y)%3,现在要求f(x,y)=i(0<=i<=2)的路径

思路：

两种解法，一种树形dp，一种点分治。

树形dp：二次扫描换根，第一次扫描，搜索出每个节点到其子树的答案贡献，第二次扫描，搜出这个点到非子树节点的答案贡献

最后统计答案的时候加上两次扫描每个点的贡献

代码：

#include<bits/stdc++.h>
using namespace std;
#define inf 0x3f3f3f3f
#define ll long long
const int maxn=1e4+9;
const int mod=1e9+7;
ll dp[maxn][3],ans[3],head[maxn],vis[maxn],fdp[maxn][3],ct[maxn][3],fct[maxn][3];
struct Edge{
    int next,to,val;
}edge[maxn*2];
int cnt=0;
void add(int u,int v,int val){
    edge[cnt].next=head[u];
    edge[cnt].to=v;
    edge[cnt].val=val;
    head[u]=cnt++;
}
void dfs1(int u){
    ct[u][0]=1;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].to;
        if(!vis[v]){
//            cout<<u<<' '<<v<<endl;
            vis[v]=1;
            dfs1(v);
            int dis=edge[i].val;
            for(int j=0;j<3;j++){
//                if(dp[v][(j-(dis%3)+3)%3]>=0){
                    ct[u][j]=(ct[u][j]+ct[v][(j-(dis%3)+3)%3])%mod;
                    dp[u][j]=(dp[u][j]+dp[v][(j-(dis%3)+3)%3]+dis*ct[v][(j-(dis%3)+3)%3])%mod;
                }
//            else
//                if(j==dis%3){
//                    ct[u][j]=max(ct[u][j],0)+1;
//                    dp[u][j]=max(dp[u][j],0)+dis;
//                }
//            }
        }
    }
}
void dfs2(int u){
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].to;
        if(!vis[v]){
            vis[v]=1;
            int dis=edge[i].val;
            for(int j=0;j<3;j++){
                fct[v][j]=(fct[u][((j-dis)%3+3)%3]+ct[u][((j-dis)%3+3)%3]-ct[v][((j-2*dis)%3+3)%3])%mod;
                fdp[v][j]=(fdp[u][((j-dis)%3+3)%3]+(dp[u][((j-dis)%3+3)%3]-dp[v][((j-2*dis)%3+3)%3]-dis*ct[v][((j-2*dis)%3+3)%3]%mod+mod)%mod+dis*fct[v][(j)%3]%mod)%mod;

            }
            dfs2(v);
        }
    }
}
int main(){
    int i,j,k,n;
    while(cin>>n){
        cnt=0;
        memset(head,-1,sizeof(head));
        memset(dp,0,sizeof(dp));
        memset(fdp,0,sizeof(fdp));
        memset(ct,0,sizeof(ct));
        memset(fct,0,sizeof(fct));
        for(i=1;i<n;i++){
            int x,y,z;cin>>x>>y>>z;
            add(x,y,z);add(y,x,z);
        }
        for(i=0;i<n;i++)vis[i]=0;
        vis[0]=1;
        dfs1(0);
        for(i=0;i<n;i++)vis[i]=0;
        vis[0]=1;
        dfs2(0);
        memset(ans,0,sizeof(ans));
        for(i=0;i<n;i++){
//            for(j=0;j<3;j++)cout<<i<<' '<<dp[i][j]<<endl;
            for(j=0;j<3;j++)(ans[j]+=fdp[i][j]+dp[i][j])%=mod;
        }
        for(i=0;i<3;i++){
            if(i!=2)cout<<ans[i]<<' ';
            else cout<<ans[i]<<endl;
        }
    }
}

点分治：点分治板子一套，稍微改改就行

代码：

#include<bits/stdc++.h>
using namespace std;
#define inf 0x3f3f3f3f
#define ll long long
const int maxn=1e4+9;
const int mod=1e9+7;
struct Edge{
    int next,to,val;
}edge[maxn*2];
int cnt=0,vis[maxn],root,k,mx=inf,Size,head[maxn],sz[maxn],r=0;
ll d[maxn],tmp[3],ans[3],id_path[3],id_num[3];
inline void add(int u,int v,int val){
    edge[cnt].next=head[u];
    edge[cnt].to=v;
    edge[cnt].val=val;
    head[u]=cnt++;
}
void getroot(int u,int fa){
    sz[u]=1;int num=0;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].to;
        if(v!=fa&&!vis[v]){
            getroot(v,u);
            sz[u]+=sz[v];
            num=max(num,sz[v]);
        }
    }
    num=max(num,Size-sz[u]);
    if(num<mx)mx=num,root=u;
}
void getdis(int u,int fa){
    id_path[d[u]%3]++;
    (id_num[d[u]%3]+=d[u])%=mod;
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].to;
        if(v!=fa&&!vis[v]){
            d[v]=(edge[i].val+d[u])%mod;
            getdis(v,u);
        }
    }
}
void cal(int u,int val){
    d[u]=val;
    for(int i=0;i<3;i++){
        id_path[i]=0;id_num[i]=0;
        tmp[i]=0;
    }
    getdis(u,0);
    for(int i=0;i<3;i++){
        for(int j=0;j<3;j++){
            tmp[(i+j)%3]=(tmp[(i+j)%3]+id_num[i]*id_path[j]+id_num[j]*id_path[i])%mod;
        }
    }
}
void dfs(int u){
    vis[u]=1;
    cal(u,0);
    for(int i=0;i<3;i++){
        ans[i]=(ans[i]+tmp[i])%mod;
    }
    for(int i=head[u];i!=-1;i=edge[i].next){
        int v=edge[i].to;
        if(!vis[v]){
            cal(v,edge[i].val);
            for(int j=0;j<3;j++){
                ans[j]=(ans[j]-tmp[j]+mod)%mod;
            }
            Size=sz[v];
            mx=inf;
            getroot(v,0);
            dfs(root);
        }
    }
}
int main(){
    int i,j,n;
    while(cin>>n){
        cnt=0;
        for(i=0;i<n;i++){
            head[i]=-1;
            vis[i]=0;
    
        }
        memset(ans,0,sizeof(ans));
        for(i=0;i<n-1;i++){
            int x,y,z;cin>>x>>y>>z;
            add(x,y,z);add(y,x,z);
        }
        Size=n;
        mx=inf;
        getroot(0,0);
        dfs(0);
        for(i=0;i<3;i++){
            if(i!=2)
            cout<<ans[i]<<' ';
            else{
                cout<<ans[i]<<endl;
            }
        }
    }
}