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);

    }

}