POJ - 1741 Tree 【树的点分治模板题】

//树的点分治，该代码是求树上有多少对点的dis<=k
//vis代表该点是否已经当做过重心，siz是子树节点个数、mv是子树中最大的节点数
int n, cnt, head[maxn], k, vis[maxn], root, maxx, dis[maxn];
int ans, num, tot, siz[maxn], mv[maxn]; //tot代表当前树的节点个数、如果遇到时间过长，考虑重心是否找对
struct node {
int to, w, next;
} e[maxn<<1];

void add(int u, int v, int w)
{
e[cnt] = (node){v,w,head[u]};
head[u] = cnt++;
}

void getroot(int u, int fa)
{
siz[u] = 1, mv[u] = 0;
for (int i = head[u]; ~i; i = e[i].next) {
int to = e[i].to;
if (to == fa || vis[to]) continue;
getroot(to, u);
siz[u] += siz[to];
mv[u] = max(mv[u], siz[to]);
}
mv[u] = max(mv[u], tot - siz[u]);
if (mv[u] < mv[root]) root = u;
}

void getdis(int u,int fa,int dep)
{
dis[++num] = dep;
for (int i = head[u]; ~i; i = e[i].next) {
int to = e[i].to;
if (to == fa || vis[to]) continue;
getdis(to, u, dep + e[i].w);
}
}

int cal(int u,int f)
{
int res = 0;
num = 0;
getdis(u,-1,f);
sort(dis+1,dis+num+1);
/*if(u == 1) {
cout << "###" << num << endl;
for(int i=1;i<=num;i++){
printf("%d%c",dis[i],i==num?'\n':' ' );
}
}*/
int r = num ;
for(int l = 1; l < r; l++) {
while(dis[l] + dis[r] > k && l < r)
r--;
res += r - l ;
}
return res;
}

void work(int u)
{
vis[u] = 1;
ans += cal(u, 0);
//if(u == 1) cout << "!!!" << ans << endl;
for (int i = head[u]; ~i; i = e[i].next) {
int to = e[i].to;
if (vis[to]) continue;
int tmp = cal(to, e[i].w);
//传e[i].w是因为这是根节点到该点的距离,因为要和k判断关系,所以必须加上
//if(u == 1) cout << "@@@" << tmp << endl;
ans -= tmp;
mv[root=0] = tot = siz[to];
getroot(to, -1);
work(root);
}
}

void solve()
{
while(~scanf("%d%d",&n,&k)){
if (n + k == 0 ) break;
cnt =0 ; Fill(head,-1); Fill(vis,0);
for (int i = 1; i < n; i++) {
int u, v, w;
scanf("%d%d%d",&u,&v,&w);
add(u, v, w); add(v, u, w);
}
ans = 0 ;
mv[root=0] = tot = n;
getroot(1, -1);
work(root);
cout << ans << endl;
}
}