POJ1741 Tree 点分治

问题描述

POJ1741

题解

题意是求树上满足两点之间距离 \(dis \le k\) 的点对 \((x,y)\) 的数目。

点分治。

\(\mathrm{Code}\)

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>

using namespace std;

template <typename Tp>
void read(Tp &x){
    x=0;char ch=1;int fh;
    while(ch!='-'&&(ch>'9'||ch<'0')) ch=getchar();
    if(ch=='-') ch=getchar(),fh=-1;
    else fh=1;
    while(ch>='0'&&ch<='9') x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
    x*=fh;
}

const int maxn=10007;
const int maxm=20007;
const int INF=0x3f3f3f3f;

int n,k;
int Head[maxn],to[maxm],Next[maxm],w[maxm],tot;
int size[maxn],mx,sz,root;

int ans,dis[maxn];
int q[maxn],l,r;

bool del[maxn];

void reset(){
    memset(Head,0,sizeof(Head));
    memset(Next,0,sizeof(Next));
    memset(del,0,sizeof(del));
    tot=ans=0;
}

void add(int x,int y,int z){
    to[++tot]=y,Next[tot]=Head[x],Head[x]=tot,w[tot]=z;
}

void Init(){
    read(k);
    for(int i=1,x,y,z;i<n;i++){
        read(x);read(y);read(z);
        add(x,y,z);add(y,x,z);
    }
}

void getdis(int x,int fa){
    q[++r]=dis[x];
    for(int i=Head[x];i;i=Next[i]){
        int y=to[i];
        if(y==fa||del[y]) continue;
        dis[y]=dis[x]+w[i];
        getdis(y,x);
    }
}

int calc(int x,int val){
    r=0;dis[x]=val;
    getdis(x,0);
    int res=0;l=1;
    sort(q+1,q+r+1);
    while(l<r){
        if(q[l]+q[r]<=k) res+=r-l,++l;
        else --r;
    }
    return res;
}

void getroot(int x,int fa){
    size[x]=1;int num=0;
    for(int i=Head[x];i;i=Next[i]){
        int y=to[i];
        if(y==fa||del[y]) continue;
        getroot(y,x);
        size[x]+=size[y];
        num=max(num,size[y]);
    }
    num=max(num,sz-size[x]);
    if(num<mx) mx=num,root=x;
}

void dfs(int x){
    ans+=calc(x,0);del[x]=1;
    for(int i=Head[x];i;i=Next[i]){
        int y=to[i];
        if(del[y]) continue;
        ans-=calc(y,w[i]);
        sz=size[y],mx=INF;
        getroot(y,0);dfs(root);
    }
}

void solve(){
    sz=n,mx=INF;
    getroot(1,0);
    dfs(root);
    printf("%d\n",ans);
}

int main(){
    while(1){
        read(n);if(!n) return 0;
        reset();Init();solve();
    }
    return 0;
}