poj1741Tree

链接：http://poj.org/problem?id=1741

题意：给定一棵有边权的树，求有多少点对满足dis(u,v)<=k。

分析：树分治裸题。看论文吧。

代码：

#include<map>
#include<set>
#include<cmath>
#include<queue>
#include<bitset>
#include<math.h>
#include<vector>
#include<string>
#include<stdio.h>
#include<cstring>
#include<iostream>
#include<algorithm>
#pragma comment(linker, "/STACK:102400000,102400000")
using namespace std;
const int N=10010;
const int mod=1000000007;
const int MOD1=1000000007;
const int MOD2=1000000009;
const double EPS=0.00000001;
typedef long long ll;
const ll MOD=1000000007;
const int INF=1000000010;
const ll MAX=1ll<<55;
const double eps=1e-5;
const double inf=~0u>>1;
const double pi=acos(-1.0);
typedef double db;
typedef long double ldb;
typedef unsigned int uint;
typedef unsigned long long ull;
int n,k,sum,ans,root;
int f[N],vis[N],siz[N],dis[N],deep[N];
int tot,u[N],v[2*N],w[2*N],pre[2*N];
void add(int a,int b,int c) {
    v[tot]=b;w[tot]=c;pre[tot]=u[a];u[a]=tot++;
    v[tot]=a;w[tot]=c;pre[tot]=u[b];u[b]=tot++;
}
void init() {
    int i,a,b,c;f[0]=N;
    ans=tot=root=0;sum=n;
    memset(u,-1,sizeof(u));
    memset(vis,0,sizeof(vis));
    for (i=1;i<n;i++)
    scanf("%d%d%d", &a, &b, &c),add(a,b,c);
}
void getroot(int x,int fa) {
    siz[x]=1;f[x]=0;
    for (int i=u[x];i!=-1;i=pre[i])
    if (vis[v[i]]||v[i]==fa) continue ;
    else {
        getroot(v[i],x);
        siz[x]+=siz[v[i]];
        f[x]=max(f[x],siz[v[i]]);
    }
    f[x]=max(f[x],sum-siz[x]);
    if (f[x]<f[root]) root=x;
}
void getdeep(int x,int fa) {
    deep[++deep[0]]=dis[x];
    for (int i=u[x];i!=-1;i=pre[i])
    if (vis[v[i]]||v[i]==fa) continue ;
    else dis[v[i]]=dis[x]+w[i],getdeep(v[i],x);
}
int cala(int x,int now) {
    dis[x]=now;deep[0]=0;
    getdeep(x,root);
    sort(deep+1,deep+deep[0]+1);
    int ret=0,l=1,r=deep[0];
    while (l<r)
    if (deep[l]+deep[r]<=k) ret+=r-l,l++;
    else r--;
    return ret;
}
void solve(int x) {
    ans+=cala(x,0);vis[x]=1;
    for (int i=u[x];i!=-1;i=pre[i])
    if (vis[v[i]]) continue ;
    else {
        ans-=cala(v[i],w[i]);
        sum=siz[v[i]];root=0;
        getroot(v[i],x);solve(root);
    }
}
int main()
{
    while (scanf("%d%d", &n, &k)&&(n||k)) {
        init();
        getroot(1,root);
        solve(root);
        printf("%d\n", ans);
    }
    return 0;
}