POJ 1741 Tree 点分治

POJ 1741

一道点分治模板题 树上距离小于等于k 实际上由于是最短距离 我们在统计上级子树的时候容易把下面子树内的情况考虑进去 导致多选 所以要给下面子树加爸爸边的贡献 删除这些影响 这点也是点分治的精华所在

/*
POJ 1741
*/
#include <cstdio>
#include <cstring>
#include <queue>
#include <cmath>
#include <iostream>
#include <stack>
#include <set>
#include <map>
#include <sstream>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;
const int MAX_N = 10024;
int f[MAX_N],siz[MAX_N];
struct edge{
    int next,v,dis;
}e[MAX_N<<1];
int eid,p[MAX_N],use[MAX_N],Siz,rt,cnt,d[MAX_N],k,ans;
inline int read()
{
    int date = 0,m = 1;
    char ch = 0;
    while(ch!='-'&&(ch<'0'||ch>'9'))
        ch = getchar();
    if(ch=='-')
    {
        m = -1;
        ch = getchar();
    }
    while(ch>='0' && ch<='9')
    {
        date = date*10+ch-'0';
        ch = getchar();
    }
    return date*m;
}
inline void write(ll qw)
{
    if(qw<0)
    {
        putchar('-');
        qw = -qw;
    }
    if(qw>9)
        write(qw/10);
    putchar(qw%10+'0');
}
void init(){
    memset(p,-1,sizeof(p));
    memset(use,0,sizeof(use));
    memset(d,0,sizeof(d));
    rt = 0;
    eid = 0;
}
void Insert(int u,int v,int dis){
    e[eid].v = v;
    e[eid].dis = dis;
    e[eid].next = p[u];
    p[u] = eid++;
}
void get_rt(int u,int fa){//u为当前点，fa为父亲节点
    f[u] = 0, siz[u] = 1;//f表示这个点最大子树的大小，siz是这个点子树大小的和
    for(int i = p[u];i!=-1;i = e[i].next){//枚举儿子
        int y = e[i].v;
        if(use[y]||y==fa) continue;//use表示之前遍历过了，这里没啥用
        get_rt(y,u);//往下遍历
        f[u] = max(f[u],siz[y]);//更新f
        siz[u] += siz[y];
    }
    f[u] = max(f[u],Siz - siz[u]);//Siz表示在现在这棵子树中点的总数，开始时Siz=n，除了枚举的儿子所在的子树外，还有一棵子树是上面的那一堆，容斥原理
    if(f[u]<f[rt]) rt = u;//更新root
}

int look1(int l,int x){
    int ans = 0,r = cnt;
    while(l<=r){
        int mid(l+r>>1);
        if(d[mid]<x) l = mid + 1;
        else ans = mid, r = mid - 1;
    }
    return ans;
}

int look2(int l,int x){
    int ans = 0,r = cnt;
    while(l<=r){
        int mid(l+r>>1);
        if(d[mid]<=x) ans = mid,l = mid + 1;
        else r = mid - 1;
    }
    return ans;
}

void get_num(int u,int fa,int D){
    for(int i = p[u];i!=-1;i=e[i].next){
        int y = e[i].v;
        if(use[y]||y==fa) continue;
        d[++cnt] = D + e[i].dis;
        get_num(y,u,d[cnt]);
    }
}

int get_ans(int u,int D){//solve
    d[cnt=1] = D;
    get_num(u,0,D);
    sort(d+1,d+1+cnt);
    /*int l = 1,ans = 0;
    while(l<cnt&&d[l]+d[cnt]<k) ++l;
    while(l<cnt&&k-d[l]>=d[l]){
        int D1 = look1(l+1,k-d[l]);
        int D2 = look2(l+1,k-d[l]);
        if(D2>=D1) ans+=D2-D1+1;
        ++l;
    }
    return ans;*///等于k

    int l = 1,r = cnt,tep = 0;
    while(l<=r){
        if(d[r]+d[l]<=k) {tep = tep + r - l;l++;}
        else r--;
    }
    return tep;//小于等于k

}

void dfs(int u){//Divide
    use[u] = 1,ans+=get_ans(u,0);
    for(int i = p[u];i!=-1;i=e[i].next){
        int y = e[i].v;
        if(use[y]) continue;
        ans-=get_ans(y,e[i].dis);
        Siz = siz[y],rt = 0;
        get_rt(y,u),dfs(rt);
    }
    return ;
}
int main(){
    int n;
    while(scanf("%d%d",&n,&k)==2){
        if(n==0) break;
        cnt = 1;
        init();
        for(int i = 1;i<n;++i){
            int a,b,dis;
            a = read(),b = read(),dis = read();
            Insert(a,b,dis);
            Insert(b,a,dis);
        }
        ans = 0;
        f[rt] = Siz = n;
        get_rt(1,0);
        dfs(1);
        printf("%d\n",ans);
    }
    return 0;
}