poj 1741 Tree 树的分治
题意：给定一棵边带权的树，求多少对点之间的距离小于k
1 #include<iostream>
2 #include<cmath>
3 #include<cstdio>
4 #include<cstring>
5 #include<algorithm>
6 using namespace std;
7 #define NMAX 10010
8 struct node
9 {
10 int num,weight;
11 node *next;
12 };
13 node *graph[NMAX];
14 node memo[2*NMAX];
15 bool use[NMAX];
16 int size[NMAX],d[NMAX];
17 int f[NMAX];
18 int n,m,top,tot,root;
19 int sumsize,minsize;
20 void add(int x,int y,int z)
21 {
22 node *p=&memo[top++];
23 p->num=y; p->weight=z; p->next=graph[x]; graph[x]=p;
24 p=&memo[top++];
25 p->num=x; p->weight=z; p->next=graph[y]; graph[y]=p;
26 }
27 void get_root(int i,int fa)
28 {
29 int big=-1;
30 size[i]=1;
31 for(node *p=graph[i];p;p=p->next)
32 if(!use[p->num]&&p->num!=fa)
33 {
34 get_root(p->num,i);
35 size[i]+=size[p->num];
36 if(size[p->num]>big) big=size[p->num];
37 }
38 if(sumsize-size[i]>big) big=sumsize-size[i];
39 if(big<minsize) minsize=big,root=i;
40 }
41 void get_dis(int i,int dis,int fa)
42 {
43 d[++tot]=dis;
44 for(node *p=graph[i];p;p=p->next)
45 if(!use[p->num]&&p->num!=fa)
46 get_dis(p->num,dis+p->weight,i);
47 }
48
49 int count(int x,int dis)
50 {
51 int s=0;
52 tot=0;
53 get_dis(x,dis,-1);
54 sort(d+1,d+tot+1);
55 for(int i=1,j=tot;i<=j;i++)
56 {
57 while(d[i]+d[j]>m&&i<=j) j--;
58 if(i<j) s+=j-i;
59 }
60 return s;
61 }
62
63 void dfs(int i)
64 {
65 sumsize=minsize=size[i];
66 get_root(i,-1);
67 i=root;
68 f[i]=count(i,0);
69 //cout<<i<<" w "<<f[i]<<endl;
70 use[i]=1;
71 for(node *p=graph[i];p;p=p->next)
72 if(!use[p->num])
73 {
74 f[i]-=count(p->num,p->weight);
75 dfs(p->num);
76 }
77 }
78 int main()
79 {
80 int i;
81 int x,y,z;
82 while(scanf("%d%d",&n,&m)!=EOF&&n)
83 {
84 int ans=0;
85 top=0;
86 memset(graph,0,sizeof(graph));
87 memset(use,0,sizeof(use));
88 for(i=1;i<n;i++)
89 {
90 scanf("%d%d%d",&x,&y,&z);
91 add(x,y,z);
92 }
93 size[1]=n;
94 dfs(1);
95 for(i=1;i<=n;i++) ans+=f[i];
96 printf("%d\n",ans);
97 }
98 return 0;
99 }
