poj1741 Tree

树分治。

网上有部分人的题解写的啥呀。。。

按重心进行分治。

首先O(n)算出以当前点为根depth(u)+depth(v)<=k的(u,v)的个数。

然后再减去同一个子树里depth(u)+depth(v)<=k的(u,v)的个数。

#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
const int maxn = 10000 + 10;
const int maxm = 20000 + 10;
const int inf = 0x3f3f3f3f;
int g[maxn],v[maxm],next[maxm],c[maxm],eid;
int n,k,tot,root,cnt,res;
int size[maxn],maxv[maxn];
bool vis[maxn];
int dis[maxn];

void addedge(int a,int b,int C) {
    v[eid]=b; c[eid]=C; next[eid]=g[a]; g[a]=eid++;
    v[eid]=a; c[eid]=C; next[eid]=g[b]; g[b]=eid++;    
}

void get(int u,int f) {
    size[u]=1; maxv[u]=0;
    for(int i=g[u];~i;i=next[i]) if(v[i]!=f&&!vis[v[i]]) {
        get(v[i],u);
        size[u]+=size[v[i]];
        maxv[u]=max(maxv[u],size[v[i]]);
    }
    maxv[u]=max(maxv[u],tot-size[u]);
    if(maxv[u]<maxv[root]) root=u;
}

void dfsdis(int u,int d,int f) {
    dis[++cnt]=d;
    for(int i=g[u];~i;i=next[i]) if(v[i]!=f&&!vis[v[i]]) dfsdis(v[i],d+c[i],u);
}

int calc(int u,int d) {
    int ret=0; cnt=0;
    dfsdis(u,d,0);
    sort(dis+1,dis+cnt+1);
    int i=1,j=cnt;
    while(i<j) {
        while(dis[i]+dis[j]>k&&i<j) j--;
        ret+=j-i;
        i++;
    }
    return ret;
}

void dfs(int u) {
    if(tot<=1) return;
    root=0; get(u,0); 
    int tmp=root; vis[tmp]=1;
    res+=calc(tmp,0);
    for(int i=g[tmp];~i;i=next[i]) {
        if(!vis[v[i]]) {
            res-=calc(v[i],c[i]);
            tot=size[v[i]]; dfs(v[i]);    
        }
    }
}

int main() {
    maxv[0]=inf;
    while(scanf("%d%d",&n,&k)==2) {
        if(!n&&!k) break;
        memset(g,-1,sizeof(g)); eid=0;
        memset(vis,0,sizeof(vis)); res=0;
        for(int i=1,u,v,w;i<n;i++) {
            scanf("%d%d%d",&u,&v,&w);
            addedge(u,v,w);
        }
        tot=n; dfs(1);
        printf("%d\n",res);    
    }
    return 0;
}