P4178 Tree (点分治)

Posted 2021-05-21 昵称很长很长真是太好了

题意：
给定一棵 n 个节点的树，每条边有边权，求出树上两点距离小于等于 k 的点对数量。
题解：
根点分治模板提很相似，只不过这个题目让你去维护小于等于k距离点的个数，这个时候我们还是要用到桶的思想，只不过每次查询查的时小于等于(k-某个值)的个数，对于这个操作，我们可以用一个树状数组来维护即可。

做的时候还是很多方面的小bug，比如说找重心的时候手残…等等

代码：

#include<bits/stdc++.h>

using namespace std;

const int maxn=1e5+10;

//int n,m;
int c[20005],k; //对应原数组和树状数组

int lowbit(int x){
    return x&(-x);
}

void updata(int i,int k){    //在i位置加上k
    while(i <= 20005){
        c[i] += k;
        i += lowbit(i);
    }
}
int getsum(int i){        //求A[1 - i]的和  y-（x-1）
    int res = 0;
    while(i > 0){
        res += c[i];
        i -= lowbit(i);
    }
    return res;
}

struct E{
    int to,next,w;
}edge[maxn];
int tot,head[maxn];
bool visited[maxn];
int maxp[maxn],sz[maxn],dis[maxn],temp[maxn];
int sum,rt,ans,n,cnt;
void add(int u,int v,int w){
    edge[tot].to=v;
    edge[tot].w=w;
    edge[tot].next=head[u];
    head[u]=tot++;
}

void getrt(int x,int fa){
    sz[x]=1,maxp[x]=0;
    for(int i=head[x];~i;i=edge[i].next){
        int v=edge[i].to;
        if(v==fa||visited[v]) continue;
        getrt(v,x);
        sz[x]+=sz[v];
        if(sz[v]>maxp[x]) maxp[x]=sz[v];
    }
    maxp[x]=max(maxp[x],sum-sz[x]);
    if(maxp[x]<maxp[rt]) rt=x;
}

void getdis(int x,int fa){
    temp[cnt++]=dis[x];
    //updata(dis[x],1);
    if(dis[x]<=k) ans++;
    for(int i=head[x];~i;i=edge[i].next){
        int v=edge[i].to;
        if(v==fa||visited[v]) continue;
        dis[v]=dis[x]+edge[i].w;
        getdis(v,x);
    }
}

void sol(int x){
    queue<int> que;
    for(int i=head[x];~i;i=edge[i].next){
        int v=edge[i].to;
        if(visited[v]) continue;
        dis[v]=edge[i].w;
        cnt=0;
        getdis(v,x);
//        for(int j=0;j<cnt;j++) cout<<temp[j]<<" ";
//        cout<<endl;
        for(int j=0;j<cnt;j++){
            ans+=getsum(k-temp[j]);
        }
        for(int j=0;j<cnt;j++){
            que.push(temp[j]);
            updata(temp[j],1);
            //cout<<"sum   "<<getsum(temp[j])<<endl;
        }
    }
    while(!que.empty()){
        updata(que.front(),-1);
        que.pop();
    }
}

void divide(int x){
    visited[x]=true;
    sol(x);
    for(int i=head[x];~i;i=edge[i].next){
        int v=edge[i].to;
        if(visited[v]) continue;
        maxp[rt=0]=sum=sz[v];
        getrt(v,0);
        getrt(rt,0);
        divide(rt);
    }
}


int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    memset(head,-1,sizeof head);
    cin>>n;
    for(int i=1;i<n;i++){
        int x,y,z;
        cin>>x>>y>>z;
        add(x,y,z);
        add(y,x,z);
    }
    cin>>k;
    maxp[0]=sum=n;
    getrt(1,0);
    //cout<<rt<<endl;
    getrt(rt,0);
    divide(rt);
    cout<<ans<<endl;
}