POJ1741 Tree（点分治）

题目

给一棵边带权树，问两点之间<=K的点对有多少个。

思路

题目很简单，但是思路很经典。

首先确定点分治的基本框架，假设一定要经过一个根。

下面还要用到容斥的思维。

对于一个根，我们没法直接统计路径长度不超过k的路径条数，那需要一点技巧。

处理出子树中所有的dis值放入B数组中，再对于当前A子树，加上其对于B所做的贡献，减去其对自身所做的贡献，就是它对其他子树所做的贡献。

看似简单，其实在点分治问题中这是非常常用的技巧。

代码

#include<stdio.h>
#include<algorithm>
#include<string.h>
#include<iostream>
#define M 10005
#define clr(x,y) memset(x,y,sizeof(x))
using namespace std;
int n,K,h[M],tot;
struct edge{
    int nxt,to,co;
}G[M<<1];
void Add(int a,int b,int c){
    ++tot;
    G[tot].nxt=h[a],G[tot].to=b,G[tot].co=c;
    h[a]=tot;
}
int sz[M],mx=2e9,zx,tt,ans;
bool vis[M];
void dfs_zx(int x,int f){
    sz[x]=1;int mm=0;
    for(int i=h[x];i;i=G[i].nxt){
        int u=G[i].to;
        if(u==f||vis[u])continue;
        dfs_zx(u,x);
        sz[x]+=sz[u];
        if(sz[u]>mm)mm=sz[u];
    }
    mm=max(mm,tt-sz[x]);
    if(mm<mx)mx=mm,zx=x;
}
void find_zx(int x){
    mx=2e9;
    dfs_zx(x,0);
}
int A[M],top,B[M],top2;
void dfs(int x,int f,int ds){
    A[++top]=ds;
    for(int i=h[x];i;i=G[i].nxt){
        int u=G[i].to,v=G[i].co;
        if(vis[u]||u==f)continue;
        dfs(u,x,ds+v);
    }
}
int find(int *s,int l,int r,int x){
    s[r]=2e9;
    int as=-1;
    while(l<=r){
        int mid=(l+r)>>1;
        if(s[mid]>x){
            as=mid;
            r=mid-1;    
        }
        else l=mid+1;
    }
    if(as==-1)as=1;
    return as-1;
}
void solve(int x){
    find_zx(x);
    x=zx;
    vis[x]=1;
    top=0;
    dfs(x,0,0);top2=0;
    for(int i=2;i<=top;i++)B[++top2]=A[i];
    sort(B+1,B+top2+1);
    for(int i=h[x];i;i=G[i].nxt){
        int u=G[i].to,v=G[i].co;
        if(vis[u])continue;
        top=0;dfs(u,x,v);
        sort(A+1,A+top+1);
        for(int j=1;j<=top;j++)ans+=find(B,1,top2+1,K-A[j])-find(A,1,top+1,K-A[j]); 
        ans+=find(A,1,top+1,K)*2;
    }
    for(int i=h[x];i;i=G[i].nxt){
        int u=G[i].to;
        if(vis[u])continue;
        tt=sz[u];
        solve(u);
    }
}
int main(){
    while(scanf("%d%d",&n,&K),n+K){
        for(int i=1;i<=n;i++)h[i]=vis[i]=0;     
        tot=ans=0;tt=n;
        for(int i=1,a,b,c;i<n;i++){
            scanf("%d%d%d",&a,&b,&c);
            Add(a,b,c);Add(b,a,c);
        }
        solve(1);
        printf("%d\n",ans/2);
    }
    return 0;   
}