「hdu5314」Happy King【点分治+树状数组】

题目链接

• http://acm.hdu.edu.cn/showproblem.php?pid=5314

题意

• 给你一颗树，求有多少有序点对$(u,v)$使得路径上所有点权最大值减去最小值不大于$D$

题解

• 挺裸的点分治
• 注意题意是有序点对，而不是加上每个点和自己构成的路径
• 每次找出当前联通块的重心，求出所有通过重心的方案数，去掉同一棵树的两个点对应的方案，求的时候可以按照最大值升序排序，然后用树状数组去查一下区间和就行了（类似树状数组求逆序对的做法），当然也可以给最小值升序排序，然后前缀二分查询当前最大值-$D$，这种方法慎用

复杂度

• $O(n(\log n{)}^2)$

代码

#pragma comment(linker, "/STACK:102400000,102400000") 
#include<cstdio>
#include<iostream>
#include<algorithm>
using namespace std;
const int maxn=2e5+10;
const int maxm=3e6+10;
int n,k,a[maxn],c[maxn],d[maxn],m;
namespace bit{
    int s[maxn];
    int lowbit(int x) {return x&(-x);} 
    void add(int id,int val){
        for(int i=id;i<=m;i+=lowbit(i)) s[i]+=val; //注意这里的n,需要根据题目要求改
    }
    int query(int id){
        int ans=0;
        for(int i=id;i>=1;i-=lowbit(i)) ans+=s[i];
        return ans;
    }
    int sum(int l,int r) {return query(r)-query(l-1);}
    void init(int n) {for(int i=0;i<=n;i++) s[i]=0;}
}
using namespace bit;

namespace point_divide_and_conquer{
    int tot,head[maxn],siz[maxn],root,min_son,num; //min_son记录当前最小的 最大儿子size
    bool vis[maxn];
    struct ed{int v,w,next;}edge[2*maxn];
    struct da{
        int minn,maxx;
        da(int a=0,int b=0) {minn=a;maxx=b;}
        friend bool operator<(const da &a,const da &b) {
            return a.maxx<b.maxx;
        }
    }data[maxn];
    inline void init(int n) {
        tot=0;
        for(int i=1;i<=n;i++) head[i]=0,vis[i]=false;
    }
    inline void add_edge(int u,int v,int w) {
        edge[++tot]=ed{v,w,head[u]};
        head[u]=tot;
    }
    inline void dfs_size(int cur,int fa) {
        siz[cur]=1;
        for(int i=head[cur];i;i=edge[i].next) {
            if(edge[i].v!=fa && !vis[edge[i].v]) {
                dfs_size(edge[i].v,cur);
                siz[cur]+=siz[edge[i].v];
            }
        }
    }
    inline void dfs_root(int cur,int fa,int all) {  //all表示这个联通块的大小
        int max_son=all-siz[cur];
        for(int i=head[cur];i;i=edge[i].next) {
            if(edge[i].v!=fa && !vis[edge[i].v]) {
                max_son=max(max_son,siz[edge[i].v]);
                dfs_root(edge[i].v,cur,all);
            }
        }
        if(max_son<min_son) min_son=max_son,root=cur;
    }
    inline void dfs_roote(int cur,int fa,int minn,int maxx) {
        if(maxx-minn<=k)data[++num]=da{minn,maxx};
        for(int i=head[cur];i;i=edge[i].next) {
            if(edge[i].v!=fa && !vis[edge[i].v]) {
                dfs_roote(edge[i].v,cur,min(a[edge[i].v],minn),max(maxx,a[edge[i].v]));
            }
        }
    }
    inline long long calc(int cur,int fa,int minn,int maxx) {
        num=0;
        dfs_roote(cur,fa,minn,maxx);
        sort(data+1,data+num+1);
        for(int i=1;i<=num;i++) d[i]=data[i].minn;
        sort(d+1,d+num+1);
        m=unique(d+1,d+num+1)-d-1;
        bit::init(m);
        long long ans=0;
        for(int i=1;i<=num;i++) {
            int loc=lower_bound(d+1,d+m+1,data[i].minn)-d;
            int pos=lower_bound(d+1,d+m+1,data[i].maxx-k)-d;
            if(pos<=m) ans+=bit::sum(pos,m);
            bit::add(loc,1);
        }
        return ans;
    }
    inline long long solve(int cur) {
        min_son=0x3f3f3f3f;
        dfs_size(cur,0);
        dfs_root(cur,0,siz[cur]);
        vis[root]=true;
        long long ans=calc(root,0,a[root],a[root]);
        for(int i=head[root];i;i=edge[i].next) {
            if(!vis[edge[i].v]) {
                ans-=calc(edge[i].v,0,min(a[root],a[edge[i].v]),max(a[root],a[edge[i].v]));
            }
        }
        for(int i=head[root];i;i=edge[i].next) if(!vis[edge[i].v]) ans+=solve(edge[i].v);
        return ans;
    }
}
using namespace point_divide_and_conquer;
int main() {
    int t;scanf("%d",&t);
    while(t--) {
        scanf("%d %d",&n,&k);
        point_divide_and_conquer::init(n);
        for(int i=1;i<=n;i++) scanf("%d",&a[i]);
        for(int i=1,u,v;i<n;i++) {
            scanf("%d %d",&u,&v);
            add_edge(u,v,1);
            add_edge(v,u,1);
        }
        printf("%lld\n",2*solve(1));
    }
}