POJ 1741 Tree （树的点分治入门）

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 16172		Accepted: 5272

Description

Give a tree with n vertices,each edge has a length(positive integer less than 1001).
Define dist(u,v)=The min distance between node u and v.
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.
Write a program that will count how many pairs which are valid for a given tree.
Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l.
The last test case is followed by two zeros.
Output

For each test case output the answer on a single line.
Sample Input

5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0

Sample Output

8

view code#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <queue>
using namespace std;
const int N = 10010;
int n, k, pre
, ans, mi, rt, siz
, num;
bool vis
;

struct edge
{
int u, v, w, next;
edge() {}
edge(int u, int v, int w, int next):u(u),v(v),w(w),next(next) {}
}e[N<<1];
int ecnt;

void init()
{
ans = ecnt = 0;
memset(pre, -1, sizeof(pre));
memset(vis, 0, sizeof(vis));
}

inline void add(int u, int v, int w)
{
e[ecnt] = edge(u, v, w, pre[u]);
pre[u] = ecnt++;
}

void getroot(int u, int fa)
{
siz[u] = 1;
int mx = 0;
for(int i=pre[u]; ~i; i=e[i].next)
{
int v = e[i].v;
if(v==fa || vis[v]) continue;
getroot(v, u);
siz[u] += siz[v];
mx = max(mx, siz[v]);
}
mx = max(mx, siz[0]-siz[u]);
if(mx <mi) mi = mx, rt = u;
}

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

int calc(int u, int d)
{
int res = 0;
num = 0;
getdis(u, d, 0);
sort(dis, dis+num);
int i = 0, j = num-1;
while(i<j)// 经典。。
{
while(dis[i]+dis[j]>k && i<j) j--;
res += j-i;
i++;
}
return res;
}

void solve(int u, int cnt)
{
mi = n;
siz[0] = cnt;
getroot(u, 0);
ans += calc(rt, 0);
vis[rt] = 1;
for(int i=pre[rt]; ~i; i=e[i].next)
{
int v = e[i].v;
if(vis[v]) continue;
ans -= calc(v, e[i].w);
solve(v, siz[v]);
}

}

int main()
{
//    freopen("in.txt", "r", stdin);
while(scanf("%d%d", &n, &k)>0 && (n|k))
{
int u, v, w;
init();
for(int i=1; i<n; i++)
{
scanf("%d%d%d", &u, &v, &w);
add(u, v, w);
add(v, u, w);
}
solve(1, n);
printf("%d\n", ans);
}
return 0;
}