hdu5314 Happy King

树分治。

代码

#pragma comment(linker, "/STACK:102400000,102400000")
#include<cstdio>
#include<cstring>
#include<algorithm>
#define LL long long
#define N 300010
using namespace std;
int dp,pre[N],p[N],tt[N],cnt,q[N];
LL c[N],ans;
int n,i,a,b,w[N],vis[N],s[N],tmp,d,tot1,tot2,v[N];
struct g{
    int a,b;
}A[N],B[N];
bool cmp(g x,g y)
{
    return x.b<y.b;
}
void link(int x,int y)
{
    dp++;pre[dp]=p[x];p[x]=dp;tt[dp]=y;
}
int lowbit(int x)
{
    return x&(-x);
}
void cc(int x,int w)
{
    while (x<=n)
    {
        c[x]+=w;
        x=x+lowbit(x);
    }
}
LL getsum(int x)
{
    LL ans=0;
    while (x>0)
    {
        ans=ans+c[x];
        x=x-lowbit(x);
    }
    return ans;
}
void dfs(int x,int fa,int Min,int Max)
{
    int i;
    Min=min(Min,w[x]);
    Max=max(Max,w[x]);
    if (Max-Min<=d)
    {
    tot1++;A[tot1].a=Min;A[tot1].b=Max;
    tot2++;B[tot2]=A[tot1];
	}
	else
	return;
	i=p[x];
    while (i)
    {
        if ((!vis[tt[i]])&&(tt[i]!=fa))
            dfs(tt[i],x,Min,Max);
        i=pre[i];
    }
}
void getroot(int x,int fa,int sum)
{
    int i,flag=0;
    i=p[x];s[x]=1;
    while (i)
    {
        if ((!vis[tt[i]])&&(tt[i]!=fa))
        {
            getroot(tt[i],x,sum);
            s[x]+=s[tt[i]];
            if (s[tt[i]]>sum/2) flag=1;
        }
        i=pre[i];
    }
    if (sum-s[x]>sum/2) flag=1;
    if (!flag) tmp=x;
}
int ef(int x)
{
    int l,r,m;
    l=1;r=n;
    while (l<=r)
    {
        m=(l+r)>>1;
        if (v[m]>x)
        r=m-1;
        else
        l=m+1;
    }
    return r;
}
void work(int x,int sum)
{
    int root,i,ta,tb;
    getroot(x,0,sum);
    root=tmp;
    i=p[root];
    vis[root]=1;
    while (i)
    {
        if (!vis[tt[i]])
        {
            if (s[tt[i]]>s[root])
            work(tt[i],sum-s[root]);
            else
            work(tt[i],s[tt[i]]);
        }
        i=pre[i];
    }

    tot1=1;
    A[tot1].a=w[root];
	A[tot1].b=w[root];
    i=p[root];
    while (i)
    {
        if (!vis[tt[i]])
        {
            tot2=0;
            dfs(tt[i],0,w[root],w[root]);
            sort(B+1,B+1+tot2,cmp);
            cnt=0;
            for (int j=1;j<=tot2;j++)
            {
                ta=ef(B[j].b);
                tb=ef(B[j].b-d-1);
                ans=ans-(getsum(ta)-getsum(tb));
                ta=ef(B[j].a);cnt++;q[cnt]=ta;
                cc(ta,1);
            }

            for (int j=1;j<=cnt;j++)
            cc(q[j],-1);
        }
        i=pre[i];
    }
    sort(A+1,A+1+tot1,cmp);
    cnt=0;
    for (int j=1;j<=tot1;j++)
    {
        ta=ef(A[j].b);
        tb=ef(A[j].b-d-1);
        ans=ans+(getsum(ta)-getsum(tb));
        ta=ef(A[j].a);cnt++;q[cnt]=ta;
        cc(ta,1);
    }
    for (int j=1;j<=cnt;j++)
    cc(q[j],-1);
    vis[root]=0;
}
int main()
{
    int test;
    scanf("%d",&test);
    while (test)
    {
    test--;
    dp=0;
	memset(p,0,sizeof(p));
    scanf("%d%d",&n,&d);
    for (i=1;i<=n;i++)
    {
        scanf("%d",&w[i]);
        v[i]=w[i];
    }
    sort(v+1,v+1+n);
    for (i=1;i<=n-1;i++)
    {
        scanf("%d%d",&a,&b);
        link(a,b);
        link(b,a);
    }
    ans=0;
    work(1,n);
    printf("%I64d\n",ans*2);
    }
}