poj 1741 Tree 树形DP + 树分治 ★★★★

Tree

Time Limit: 1000MS Memory Limit: 30000KTotal Submissions: 15090 Accepted: 4913

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 41 2 31 3 11 4 23 5 10 0

Sample Output

8

Source

LouTiancheng@POJ

话说树分治好tm难啊，想了好几天，写了一天就是为了他。

然而想是想不出来的，你必须有树分治这个知识。。。

推荐题解：

http://blog.csdn.net/woshi250hua/article/details/7723400

我的代码用的是STL。

#include<cstdio>#include<string>#include<cstring>#include<iostream>#include<cmath>#include<algorithm>#include<climits>#include<queue>#include<vector>#include<map>#include<sstream>#include<set>#include<stack>#include<cctype>#include<utility>#pragma comment(linker, "/STACK:102400000,102400000")#define PI (4.0*atan(1.0))#define eps 1e-10#define sqr(x) ((x)*(x))#define FOR0(i,n)  for(int i=0 ;i<(n) ;i++)#define FOR1(i,n)  for(int i=1 ;i<=(n) ;i++)#define FORD(i,n)  for(int i=(n) ;i>=0 ;i--)#define  lson   ind<<1,le,mid#define rson    ind<<1|1,mid+1,ri#define MID   int mid=(le+ri)>>1#define zero(x)((x>0? x:-x)<1e-15)#define mk    make_pair#define _f     first#define _s     second#define ysk(x)  (1<<(x))using namespace std;//const int INF=    ;typedef long long ll;//const ll inf =1000000000000000;//1e15;//ifstream fin("input.txt");//ofstream fout("output.txt");//fin.close();//fout.close();//freopen("a.in","r",stdin);//freopen("a.out","w",stdout);const int INF =0x3f3f3f3f;const int maxn= 10000+20   ;//const int maxm=    ;vector<int >G[maxn];bool vis[maxn];int dp[maxn],maxsub[maxn],sign[maxn];struct Edge{    int from,to,w;    Edge(){}    Edge(int from,int to,int w):from(from),to(to),w(w){}};int n,k,ans,tot;vector<Edge>edges;vector<int >ve;inline void add_edge(int s,int t,int w){    edges.push_back(Edge(s,t,w));    edges.push_back(Edge(t,s,w));    int m=edges.size();    G[s].push_back(m-2);    G[t].push_back(m-1);}void init(){    for(int i=1;i<=n;i++) G[i].clear();    edges.clear();    memset(vis,0,(n+1)*sizeof vis[0]);//    memset(vis,0,sizeof vis);    ans=0;}int dfs(int x,int fa){    dp[x]=1;    maxsub[x]=0;    sign[tot++]=x;    for(int i=0;i<G[x].size();i++)    {        int id=G[x][i];        Edge & e=edges[id];        int y=e.to;        if(y==fa||vis[y])  continue;        dp[x]+=dfs(y,x);        maxsub[x]=max(maxsub[x],dp[y]);    }    return dp[x];}int findroot(int x){    tot=0;    dfs(x,-1);    int maxs=INF,p;    for(int i=0;i<tot;i++)    {        int y=sign[i];        maxsub[y]=max(maxsub[y],dp[x]-dp[y]);        if(maxsub[y]<maxs)        {            maxs=maxsub[y];            p=y;        }    }    return p;}void finddis(int x,int fa,int dis){    ve.push_back(dis);//    cout<<x<<" "<<"dis "<<dis<<endl;;    for(int i=0;i<G[x].size();i++)    {         int id=G[x][i];         Edge & e=edges[id];         int y=e.to;         int & w=e.w;         if(y==fa||vis[y]||dis+w>k )  continue;         finddis(y,x,dis+w);    }}void cal1(){    sort(ve.begin(),ve.end());    int le=0,ri=ve.size()-1;    while(le<ri)    {        if(ve[le]+ve[ri]>k)  ri--;        else  {ans+=ri-le;le++ ;  }    }}void cal2(int x){    vis[x]=1;     for(int i=0;i<G[x].size();i++)    {         int id=G[x][i];         Edge & e=edges[id];         int y=e.to;         if(vis[y])   continue;          int & w=e.w;         ve.clear();         finddis(y,x,w);          sort(ve.begin(),ve.end());          int le=0,ri=ve.size()-1;          while(le<ri)        {            if(ve[le]+ve[ri]>k)  ri--;           else  {ans-=ri-le;le++ ;  }        }    }}void solve(int x,int fa){//    cout<<"x: "<<x<<endl;    int root=findroot(x);//ok//    cout<<"root: "<<root<<endl;    ve.clear();    finddis(root,-1,0);/*    for(int i=0;i<ve.size();i++)    {        cout<<ve[i]<<" ";    }*///    cout<<endl;    cal1();//    cout<<"ans :"<<ans<<endl;    cal2(root);//    cout<<"ans :"<<ans<<endl;     for(int i=0;i<G[root].size();i++)    {         int id=G[root][i];         Edge & e=edges[id];         int y=e.to;         if(y==fa||vis[y])  continue;         solve(y,root);    }}int main(){    int x,y,w;   while(~scanf("%d%d",&n,&k)&&(n||k))   {       init();       for(int i=1;i<n;i++)       {          scanf("%d%d%d",&x,&y,&w);          add_edge(x,y,w);       }       solve(1,-1);       printf("%d\n",ans);   }    return 0;}/*2 51 2 33 51 2 42 3 23 51 2 42 3 17 51 2 12 3 41 4 74 5 21 7 56 7 31 2 3*/

#include<cstdio>#include<string>#include<cstring>#include<iostream>#include<cmath>#include<algorithm>#include<climits>#include<queue>#include<vector>#include<map>#include<sstream>#include<set>#include<stack>#include<cctype>#include<utility>#pragma comment(linker, "/STACK:102400000,102400000")#define PI (4.0*atan(1.0))#define eps 1e-10#define sqr(x) ((x)*(x))#define FOR0(i,n)  for(int i=0 ;i<(n) ;i++)#define FOR1(i,n)  for(int i=1 ;i<=(n) ;i++)#define FORD(i,n)  for(int i=(n) ;i>=0 ;i--)#define  lson   ind<<1,le,mid#define rson    ind<<1|1,mid+1,ri#define MID   int mid=(le+ri)>>1#define zero(x)((x>0? x:-x)<1e-15)#define mk    make_pair#define _f     first#define _s     second#define ysk(x)  (1<<(x))using namespace std;//const int INF=    ;typedef long long ll;//const ll inf =1000000000000000;//1e15;//ifstream fin("input.txt");//ofstream fout("output.txt");//fin.close();//fout.close();//freopen("a.in","r",stdin);//freopen("a.out","w",stdout);const int INF =0x3f3f3f3f;const int maxn= 10000+20   ;//const int maxm=    ;vector<int >G[maxn];bool vis[maxn];int dp[maxn],maxsub[maxn],sign[maxn];struct Edge{    int from,to,w;    Edge(){}    Edge(int from,int to,int w):from(from),to(to),w(w){}};int n,k,ans,tot;vector<Edge>edges;vector<int >ve;inline void add_edge(int s,int t,int w){    edges.push_back(Edge(s,t,w));    edges.push_back(Edge(t,s,w));    int m=edges.size();    G[s].push_back(m-2);    G[t].push_back(m-1);}void init(){    for(int i=1;i<=n;i++) G[i].clear();    edges.clear();    memset(vis,0,(n+1)*sizeof vis[0]);//    memset(vis,0,sizeof vis);    ans=0;}int dfs(int x,int fa){    dp[x]=1;    maxsub[x]=0;    sign[tot++]=x;    for(int i=0;i<G[x].size();i++)    {        int id=G[x][i];        Edge & e=edges[id];        int y=e.to;        if(y==fa||vis[y])  continue;        dp[x]+=dfs(y,x);        maxsub[x]=max(maxsub[x],dp[y]);    }    return dp[x];}int findroot(int x){    tot=0;    dfs(x,-1);    int maxs=INF,p;    for(int i=0;i<tot;i++)    {        int y=sign[i];        maxsub[y]=max(maxsub[y],dp[x]-dp[y]);        if(maxsub[y]<maxs)        {            maxs=maxsub[y];            p=y;        }    }    return p;}void finddis(int x,int fa,int dis){    ve.push_back(dis);//    cout<<x<<" "<<"dis "<<dis<<endl;;    for(int i=0;i<G[x].size();i++)    {         int id=G[x][i];         Edge & e=edges[id];         int y=e.to;         int & w=e.w;         if(y==fa||vis[y]||dis+w>k )  continue;         finddis(y,x,dis+w);    }}void cal1(){    sort(ve.begin(),ve.end());    int le=0,ri=ve.size()-1;    while(le<ri)    {        if(ve[le]+ve[ri]>k)  ri--;        else  {ans+=ri-le;le++ ;  }    }}void cal2(int x){//    vis[x]=1;     for(int i=0;i<G[x].size();i++)    {         int id=G[x][i];         Edge & e=edges[id];         int y=e.to;         if(vis[y])   continue;          int & w=e.w;         ve.clear();         finddis(y,x,w);          sort(ve.begin(),ve.end());          int le=0,ri=ve.size()-1;          while(le<ri)        {            if(ve[le]+ve[ri]>k)  ri--;           else  {ans-=ri-le;le++ ;  }        }    }}void solve(int x,int fa){//    cout<<"x: "<<x<<endl;    int root=findroot(x);//ok    vis[root]=1;//    cout<<"root: "<<root<<endl;    ve.clear();    finddis(root,-1,0);/*    for(int i=0;i<ve.size();i++)    {        cout<<ve[i]<<" ";    }*///    cout<<endl;    cal1();//    cout<<"ans :"<<ans<<endl;    cal2(root);//    cout<<"ans :"<<ans<<endl;     for(int i=0;i<G[root].size();i++)    {         int id=G[root][i];         Edge & e=edges[id];         int y=e.to;         if(y==fa||vis[y])  continue;         solve(y,root);    }}int main(){    int x,y,w;   while(~scanf("%d%d",&n,&k)&&(n||k))   {       init();       for(int i=1;i<n;i++)       {          scanf("%d%d%d",&x,&y,&w);          add_edge(x,y,w);       }       solve(1,-1);       printf("%d\n",ans);   }    return 0;}/*2 51 2 33 51 2 42 3 23 51 2 42 3 17 51 2 12 3 41 4 74 5 21 7 56 7 31 2 3*/