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; }
没有人不辛苦,只有人不喊疼