Tree bzoj 1468

题目

给你一棵TREE,以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于K

分析

点分治
我们找到树的重心，然后dfs，求出每个点到root的距离deep，然后对deep排序，扫描哪些点对是符合的。
但是，点分治要求处理的路径是经过root，所以如果一条路径是在同一个子树之内的就不符合要求，所以还要对子树dfs一下，然后去重。
接下来处理好root后，就可以处理其他连通块了，即递归其子树。

code

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<queue>
#include<string>
#include<cmath>

using namespace std;

const int maxn=50000+10;
const int INF=1e9;

struct arr{
    int x,y;
    int w;
    int next;
    int flag;
}edge[maxn*2];
int ls[maxn];
int edge_m;

long long ans;

int n,k;

bool done[maxn];

void add(int x,int y,int w)
{
    edge[++edge_m]=(arr){x,y,w,ls[x]},ls[x]=edge_m;
    edge[++edge_m]=(arr){y,x,w,ls[y]},ls[y]=edge_m;
}

int sz[maxn];
int f[maxn];
int rt,size;

void getrt(int x,int fa)
{
    sz[x]=1;
    f[x]=0;
    for(int i=ls[x];i;i=edge[i].next)
    {
        int u=edge[i].y;
        if(u==fa||done[u]) continue;
        getrt(u,x);
        sz[x]+=sz[u];
        f[x]=max(f[x],sz[u]);
    }
    f[x]=max(f[x],size-sz[x]);
    if(f[x]<f[rt]) rt=x;
}

int dis[maxn];  
int tot=0;

void dfs(int x,int r)
{
    sz[x]=1;
    int xx=tot;
    for (int i=ls[x];i;i=edge[i].next)
    {
        if (done[edge[i].y]) continue;
        if (edge[i].y==r) continue;
        dis[++tot]=0;
        dis[tot]=dis[xx]+edge[i].w;
        dfs(edge[i].y,x);
        sz[x]+=sz[edge[i].y];
    }
}

int cale(int x,int cost)
{
    tot=1;
    dis[tot]=cost;
    dfs(x,0);
    sort(dis+1,dis+tot+1);
    int l=1,r=tot,sum=0;
    while (l<r)
    {
        if (dis[l]+dis[r]<=k)
        {
            sum+=r-l;
            l++;
        }
        else r--;
    }
    return sum;
}

int work(int x)
{
    done[x]=1;
    ans+=cale(x,0);
    for (int i=ls[x];i;i=edge[i].next)
    {
        if (done[edge[i].y]) continue;
        ans-=cale(edge[i].y,edge[i].w);
        size=sz[edge[i].y];
        getrt(edge[i].y,rt=0);
        work(rt);
    }
}

int main()
{
    while (~scanf("%d",&n))
    {
        memset(ls,0,sizeof(ls));
        memset(edge,0,sizeof(edge));
        memset(done,0,sizeof(done));
        memset(sz,0,sizeof(sz));
        memset(f,0,sizeof(f));
        edge_m=0;

        for (int i=1;i<n;i++)
        {
            int x,y,w;
            scanf("%d%d%d",&x,&y,&w);
            add(x,y,w);
        }
        scanf("%d",&k);
        f[0]=1e9;
        ans=0;
        getrt(1,rt=0);
        work(rt);
        printf("%d",ans);
    }
}