P2634 [国家集训队]聪聪可可(点分治)

由题意可知,可以分别统计以i为根的子树距离mod3的情况,并记录个数,之后用点分治可以使复杂度降到nlogn

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=3e5+10;
const int mod=1e7+7;
int n,idx;
ll ans;
int h[N],ne[N],e[N],w[N];
int cnt[N],vis[N],d[N],sz[N];
int root;
ll gcd(ll a,ll b){
    return b?gcd(b,a%b):a;
}
void add(int a,int b,int c){
    e[idx]=b,ne[idx]=h[a],w[idx]=c,h[a]=idx++;
}
void dfs_root(int u,int fa,int tot){
    int i;
    sz[u]=1;
    int ans=0;
    for(i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa||vis[j])
            continue;
        dfs_root(j,u,tot);
        sz[u]+=sz[j];
        ans=max(ans,sz[j]);
    }
    ans=max(ans,tot-sz[u]);
    if(ans*2<=tot){
        root=u;
    }
}
void dfs_sz(int u,int fa){
    sz[u]=1;
    int i;
    for(i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa||vis[j])
            continue;
        dfs_sz(j,u);
        sz[u]+=sz[j];
    }
}
void get(int u,int fa){
    cnt[d[u]]++;
    int i;
    for(i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(vis[j]||j==fa)
            continue;
        d[j]=(d[u]+w[i])%3;
        get(j,u);
    }
}
ll cal(int u,int sum){
    d[u]=sum;
    cnt[0]=cnt[1]=cnt[2]=0;
    get(u,-1);
    return (ll)cnt[1]*cnt[2]*2+(ll)cnt[0]*cnt[0];
}
void work(int u,int tot){
    dfs_root(u,-1,tot);
    u=root;
    vis[u]=1;
    dfs_sz(u,-1);
    ans+=cal(u,0);
    for(int i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(vis[j])
            continue;
        ans-=cal(j,w[i]);
        work(j,sz[j]);
    }
}
int main(){
    ios::sync_with_stdio(false);
    cin>>n;
    int i;
    memset(h,-1,sizeof h);
    for(i=1;i<n;i++){
        int a,b,c;
        cin>>a>>b>>c;
        c%=3;
        add(a,b,c);
        add(b,a,c);
    }
    work(1,n);
    int d=gcd(ans,n*n);
    cout<<ans/d<<"/"<<n*n/d<<endl;
    return 0;
}
View Code
没有人不辛苦,只有人不喊疼
原文地址:https://www.cnblogs.com/ctyakwf/p/13542792.html