BZOJ 1468: Tree

Description

真·树,问距离不大于 \(k\) 的点对个数.

Sol

点分治.

同上.

Code

/**************************************************************
    Problem: 1468
    User: BeiYu
    Language: C++
    Result: Accepted
    Time:832 ms
    Memory:3804 kb
****************************************************************/

#include<cstdio>
#include<utility>
#include<vector>
#include<queue>
#include<algorithm>
#include<iostream>
using namespace std;

typedef pair< int,int > pr;
#define mpr make_pair
#define debug(a) cout<<#a<<"="<<a<<" "
const int N = 40005;

int n,m,k,rt,sz,ans;
vector<pr> g[N];
int s[N],f[N],b[N];
vector<int> d;
int q[N],h,t;

inline int in(int x=0,char ch=getchar()){ while(ch>'9' || ch<'0') ch=getchar();
    while(ch>='0' && ch<='9') x=(x<<3)+(x<<1)+ch-'0',ch=getchar();return x; }
void GetRoot(int u,int fa){
    s[u]=1,f[u]=0;
    for(int i=0,v;i<g[u].size();i++) if((v=g[u][i].first)!=fa && !b[v]){
        GetRoot(v,u),s[u]+=s[v],f[u]=max(f[u],s[v]);
    }f[u]=max(f[u],sz-s[u]);if(f[u]<f[rt]) rt=u;
}
void GetDeep(int u,int fa,int dep){
    d.push_back(dep),s[u]=1;
    for(int i=0,v;i<g[u].size();i++) if((v=g[u][i].first)!=fa && !b[v])
        GetDeep(v,u,dep+g[u][i].second),s[u]+=s[v];
}
int Calc(int u,int w){
    d.clear(),GetDeep(u,u,w);
    sort(d.begin(),d.end());
    int res=0;
    for(int l=0,r=d.size()-1;l<r;) if(d[l]+d[r]<=k) res+=r-l,l++;else r--;
    return res;
}
void GetAns(int u){
    ans+=Calc(u,0);b[u]=1;
    for(int i=0,v;i<g[u].size();i++) if(!b[v=g[u][i].first])
        ans-=Calc(v,g[u][i].second),f[0]=sz=s[v],GetRoot(v,rt=0),GetAns(rt);
}
int main(){
//  freopen("in.in","r",stdin);
    n=in();
    for(int i=1,u,v,w;i<n;i++) u=in(),v=in(),w=in(),g[u].push_back(mpr(v,w)),g[v].push_back(mpr(u,w));
    k=in();
    f[rt=0]=sz=n;
    GetRoot(1,0);
    GetAns(rt);
    cout<<ans<<endl;
    return 0;
}