1468: Tree

题目链接

题目大意：询问树上距离小于等于k的点对数

题解：点分治，抄抄板子就会了

我的收获：点分治T1 get

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int M=100005;
#define INF 0x3f3f3f3f

int n,k,t;
int temp,ans,tp,root;
int head[M],d[M],vis[M],sz[M],f[M],s[M];

struct edge{int to,val,nex;}e[M*4];

void add(int u,int v,int z){e[t]=edge{v,z,head[u]};head[u]=t++;}

void dfs(int x,int fa)
{
    s[++tp]=d[x];
    for(int i=head[x];i!=-1;i=e[i].nex){
        int v=e[i].to;
        if(v!=fa&&!vis[v]) d[v]=d[x]+e[i].val,dfs(v,x);
    }
}

void dfsroot(int x,int fa){
    sz[x]=1;f[x]=0;
    for(int i=head[x];i!=-1;i=e[i].nex){
        int v=e[i].to;
        if(v!=fa&&!vis[v]) dfsroot(v,x),sz[x]+=sz[v],f[x]=max(f[x],sz[v]);
    }
    f[x]=max(f[x],temp-sz[x]);
    if(f[x]<f[root]) root=x;
}

int calc(int x,int w){
    int ret=0;tp=0;
    d[x]=w,dfs(x,0);
    sort(s+1,s+tp+1);
    for(int l=1,r=tp;l<r;l++){
        while(s[l]+s[r]>k&&l<r) r--;
        ret+=r-l;
    }
    return ret;
}

void solve(int x)
{ 
    ans+=calc(x,0);vis[x]=1;
    for(int i=head[x];i!=-1;i=e[i].nex)
    if(!vis[e[i].to])  
    {
        ans-=calc(e[i].to,e[i].val);
        root=0;temp=sz[e[i].to];
        dfsroot(e[i].to,0);solve(root);
    }
}

void work()
{
    f[0]=INF,temp=n;
    dfsroot(1,0),solve(root);
    printf("%d\n",ans);
}

void init()
{
    int x,y,z;memset(head,-1,sizeof(head));
    cin>>n;t=0;
    for(int i=1;i<n;i++)
        scanf("%d%d%d",&x,&y,&z),add(x,y,z),add(y,x,z);
    cin>>k;
    dfs(1,0);
}

int main()
{
    init();
    work();
}