poj-1741（树分治 点权）

题目链接

思路：

   

   

//#pragma comment(linker, "/STACK:1024000000,1024000000")
#include<iostream>
#include<stdio.h>
#include<math.h>
#include <string>
#include<string.h>
#include<map>
#include<queue>
#include<set>
#include<utility>
#include<vector>
#include<algorithm>
#include<stdlib.h>
using namespace std;
#define eps 1e-8
#define pii pair<int,int>
#define inf 0x3f3f3f3f
#define rd(x) scanf("%d",&x)
#define rd2(x,y) scanf("%d%d",&x,&y)
#define ll long long int
#define mod 1000000007
#define maxn 11000
#define maxm 10001005
int mi(int a,int b){return a<b?a:b;}
int ma(int a,int b){return a>b?a:b;}
int n,k,mm,root;
int ans;
int vis[maxn],sz[maxn],mx[maxn];
int dis[maxn],num;
vector< pair<int,int> > vv[maxn];
void dfsize(int x,int fa){//统计子树大小
    sz[x]=1;mx[x]=0;
    for(int i=0;i<vv[x].size();i++){
        pii p=vv[x][i];
        int v=p.first;
        if(v==fa||vis[v]) continue;
        dfsize(v,x);
        sz[x]+=sz[v];
        if(sz[v]>mx[x]) mx[x]=sz[v];
    }
}
void dfsroot(int r,int x,int fa){//求重心
    if(sz[r]-sz[x]>mx[x]) mx[x]=sz[r]-sz[x];
    if(mx[x]<mm) mm=mx[x],root=x;
    for(int i=0;i<vv[x].size();i++){
        pii p=vv[x][i];
        int v=p.first;
        if(v==fa||vis[v]) continue;
        dfsroot(r,v,x);
    }
}
void dfsdis(int u,int d,int fa){//求距离
    dis[num++]=d;
    for(int i=0;i<vv[u].size();i++){
        pii p=vv[u][i];
        int v=p.first;
        if(fa==v||vis[v]) continue;
        dfsdis(v,d+p.second,u);
    }
}
int calc(int u,int d){//统计经过u的路径个数和，其中包括在同一子树路径上的
    int res=0;
    num=0;
    dfsdis(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 dfs(int u){
    mm=n;
    dfsize(u,0);
   // root=u;
    dfsroot(u,u,0);
    ans+=calc(root,0);
    vis[root]=1;
    int rt=root;
    for(int i=0;i<vv[rt].size();i++){
        pii p=vv[rt][i];
        int v=p.first;
        if(vis[v]) continue;
        ans-=calc(v,p.second);//减去belongi==belongj的个数
        dfs(v);
    }
}
int main()
{
    int u,v,z;
    while(rd2(n,k)&&(n||k)){
        for(int i=0;i<=n;i++) vv[i].clear();
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&u,&v,&z);
            vv[u].push_back(make_pair(v,z));
            vv[v].push_back(make_pair(u,z));
        }
        memset(vis,0,sizeof(vis));
        ans=0;
        dfs(1);
        printf("%d\n",ans);
    }
    return 0;
}