POJ 1741/1987 树的点分治

#include <iostream>
#include <vector>
#include <algorithm>
#include <string>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <queue>
#include <stack>
#include <map>
#include <set>
#include <cmath>
#include <ctime>
#include <cassert>
#include <sstream>
using namespace std;

const int N=10010;
struct Edge {
int to,next;
int w;
Edge(){}
Edge (int _t,int _n,int _w=0) {
to=_t;
next=_n;
w=_w;
}
}edge[N<<1];
int idx,head
;
void addEdge (int u,int v,int w) {
++idx;
edge[idx]=Edge(v,head[u],w);
head[u]=idx;
}
bool vis
;

int K;// 输入中的k

int dis
,disCnt;
void getDis(int u,int f,int d) {
dis[disCnt++]=d;
for (int k=head[u];k;k=edge[k].next) {
int v=edge[k].to;
if (vis[v]||v==f) continue;
getDis(v,u,d+edge[k].w);
}
}

int bal,cmp;
int getBal(int u,int f) {
int son=1;
int dp=0;
for (int k=head[u];k;k=edge[k].next) {
int v=edge[k].to;
if (vis[v]||v==f) continue;
int subSon=getBal(v,u);
son+=subSon;
dp=max(dp,subSon);
}
dp=max(dp,disCnt-son);
if (dp<cmp){
cmp=dp;
bal=u;
}
return son;
}

int calc(int u,int initDis) {
disCnt=0;
getDis(u,-1,initDis);
sort(dis,dis+disCnt);
int ret=0;
int l=0,r=disCnt-1;
while (l<r) {
if (dis[l]+dis[r]>K) r--;
else ret+=r-l++;
}
return ret;
}

int ans=0;
void solve(int u) {
cmp=~0U>>1;
getBal(u,-1);
ans+=calc(bal,0);
vis[bal]=true;
for (int k=head[bal];k;k=edge[k].next) {
int v=edge[k].to;
if (vis[v]) continue;
ans-=calc(v,edge[k].w);// 减去子树内的重复计算的情况
solve(v);
}
}

void init(int n) {
idx=1;memset(head,0,sizeof head);
memset(vis,false,sizeof vis);
disCnt=n;// 由于在计算重心时，可以直接用之前算dis时的disCnt，所以这里初始化为n
ans=0;
}
int main () {
int n;
while (scanf("%d %d",&n,&K)!=EOF) {
if (n==0&&K==0) break;
init(n);
for (int i=1;i<n;i++) {
int u,v,w;
scanf("%d %d %d",&u,&v,&w);
addEdge(u,v,w);
addEdge(v,u,w);
}
solve(1);
printf("%d\n",ans);
}
return 0;
}