POJ 1741 Tree （树分治入门）

Tree

Time Limit: 1000MS		Memory Limit: 30000K
Total Submissions: 8554		Accepted: 2545

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

Source

LouTiancheng@POJ

可以看论文。

讲得很明白了。

写分治主要是合并的过程。

 /* ***********************************************
 Author        :kuangbin
 Created Time  :2013-11-17 14:30:29
 File Name     :E:\2013ACM\专题学习\树的分治\POJ1741.cpp
 ************************************************ */

 #include <stdio.h>
 #include <string.h>
 #include <iostream>
 #include <algorithm>
 #include <vector>
 #include <queue>
 #include <set>
 #include <map>
 #include <string>
 #include <math.h>
 #include <stdlib.h>
 #include <time.h>
 using namespace std;
 const int MAXN = ;
 const int INF = 0x3f3f3f3f;
 struct Edge
 {
     int to,next,w;
 }edge[MAXN*];
 int head[MAXN],tot;
 void init()
 {
     tot = ;
     memset(head,-,sizeof(head));
 }
 void addedge(int u,int v,int w)
 {
     edge[tot].to = v; edge[tot].w = w;
     edge[tot].next = head[u];head[u] = tot++;
 }
 bool vis[MAXN];
 int size[MAXN],dep[MAXN];
 int le,ri;
 int dfssize(int u,int pre)
 {
     size[u] = ;
     for(int i = head[u];i != -;i = edge[i].next)
     {
         int v = edge[i].to;
         if(v == pre || vis[v])continue;
         size[u] += dfssize(v,u);
     }
     return size[u];
 }
 int minn;
 //找重心
 void getroot(int u,int pre,int totnum,int &root)
 {
     int maxx = totnum - size[u];
     for(int i = head[u];i != -;i = edge[i].next)
     {
         int v = edge[i].to;
         if(v == pre || vis[v])continue;
         getroot(v,u,totnum,root);
         maxx = max(maxx,size[v]);
     }
     if(maxx < minn){minn = maxx; root = u;}
 }
 void dfsdepth(int u,int pre,int d)
 {
     dep[ri++] = d;
     for(int i = head[u];i != -;i = edge[i].next)
     {
         int v = edge[i].to;
         if(v == pre || vis[v])continue;
         dfsdepth(v,u,d+edge[i].w);
     }
 }
 int k;
 int getdep(int a,int b)
 {
     sort(dep+a,dep+b);
     int ret = , e = b-;
     for(int i = a;i < b;i++)
     {
         if(dep[i] > k)break;
         while(e >= a && dep[e] + dep[i] > k)e--;
         ret += e - a + ;
         if(e > i)ret--;
     }
     return ret>>;
 }
 int solve(int u)
 {
     int totnum = dfssize(u,-);
     int ret = ;
     minn = INF;
     int root;
     getroot(u,-,totnum,root);
     vis[root] = true;
     for(int i = head[root];i != -;i = edge[i].next)
     {
         int v = edge[i].to;
         if(vis[v])continue;
         ret += solve(v);
     }
     le = ri = ;
     for(int i = head[root];i != -;i = edge[i].next)
     {
         int v = edge[i].to;
         if(vis[v])continue;
         dfsdepth(v,root,edge[i].w);
         ret -= getdep(le,ri);
         le = ri;
     }
     ret += getdep(,ri);
     for(int i = ;i < ri;i++)
     {
         if(dep[i] <= k)ret++;
         else break;
     }
     vis[root] = false;
     return ret;
 }

 int main()
 {
     //freopen("in.txt","r",stdin);
     //freopen("out.txt","w",stdout);
     int n;
     int u,v,w;
     while(scanf("%d%d",&n,&k) == )
     {
         if(n ==  && k == )break;
         init();
         for(int i = ;i < n;i++)
         {
             scanf("%d%d%d",&u,&v,&w);
             addedge(u,v,w);
             addedge(v,u,w);
         }
         memset(vis,false,sizeof(vis));
         printf("%d\n",solve());
     }
     return ;
 }