poj1741（入门点分治）

Posted 2022-08-07 frankchen831x

题目链接：https://vjudge.net/problem/POJ-1741

题意：给出一棵树，求出树上距离不超过k的点对数量。

思路：点分治经典题。先找重心作为树根，然后求出子树中所有点到重心的距离dis[i]，那么所有组合为dis[i]+dis[j]<=k，其中不合法组合为在重心的同一个子树内的情况，所以要减去在重心的子树中仍满足dis[i]+dis[j]<=k的情况。

AC代码：

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

const int inf=0x3f3f3f3f;
const int maxn=10005;
struct node1
    int v,w,nex;
edge[maxn<<1];

struct node2
    int x,y;
arr[maxn];

int n,k,cnt,ans,head[maxn],sz[maxn],mson[maxn],Min,root,size;
int vis[maxn],t,tt,dis[maxn],pre[maxn];

void adde(int u,int v,int w)
    edge[++cnt].v=v;
    edge[cnt].w=w;
    edge[cnt].nex=head[u];
    head[u]=cnt;


void getroot(int u,int fa)
    sz[u]=1,mson[u]=0;
    for(int i=head[u];i;i=edge[i].nex)
        int v=edge[i].v;
        if(vis[v]||v==fa) continue;
        getroot(v,u);
        sz[u]+=sz[v];
        if(sz[v]>mson[u]) mson[u]=sz[v];
    
    if(size-sz[u]>mson[u]) mson[u]=size-sz[u];
    if(mson[u]<Min) Min=mson[u],root=u;


void getdis(int u,int fa,int len)
    dis[++t]=len;
    for(int i=head[u];i;i=edge[i].nex)
        int v=edge[i].v;
        if(vis[v]||v==fa) continue;
        getdis(v,u,len+edge[i].w);
    


void solve(int x,int y,int f)
    t=0;
    getdis(x,0,y);
    sort(dis+1,dis+t+1);
    tt=0,dis[0]=-1,pre[0]=0;
    for(int i=1;i<=t;++i)
        if(dis[i]!=dis[i-1]) arr[++tt].x=dis[i],arr[tt].y=1,pre[tt]=pre[tt-1]+1;
        else ++arr[tt].y,++pre[tt];
    for(int i=1;i<=tt&&arr[i].x<=k/2;++i)
        ans+=(arr[i].y-1)*arr[i].y/2*f;
    for(int i=1;i<=tt&&arr[i].x<k/2;++i)
        int l=i+1,r=tt,mid;
        while(l<=r)
            mid=(l+r)>>1;
            if(arr[i].x+arr[mid].x<=k) l=mid+1;
            else r=mid-1;
        
        ans+=(arr[i].y)*(pre[r]-pre[i])*f;
    


void fenzhi(int u,int ssize)
    vis[u]=1;
    solve(u,0,1);
    for(int i=head[u];i;i=edge[i].nex)
        int v=edge[i].v;
        if(vis[v]) continue;
        solve(v,edge[i].w,-1);
        Min=inf,root=0;
        size=sz[v]<sz[u]?sz[v]:(ssize-sz[u]);
        getroot(v,0);
        fenzhi(root,size);
    


int main()
    while(scanf("%d%d",&n,&k),n||k)
        cnt=ans=0;
        for(int i=0;i<=n;++i)
            head[i]=0,vis[i]=0;
        for(int i=1;i<n;++i)
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            adde(u,v,w);
            adde(v,u,w);
        
        Min=inf,root=0,size=n;
        getroot(1,0);
        fenzhi(root,n);
        printf("%d\n",ans);