hdu_5314_Happy King(树的点分治)

 1 #include<bits/stdc++.h>
 2 #define F(i,a,b) for(int i=a;i<=b;++i)
 3 using namespace std;
 4 typedef pair<int,int>P;
 5 typedef long long ll;
 6 const int N=1e5+7;
 7 
 8 int n,k,g[N],v[N*2],nxt[N*2],ed,w[N],t;
 9 int vis[N],size[N],mx[N],mi,tot,root;
10 P dis[N];
11 ll ret;
12 
13 inline void adg(int x,int y){v[++ed]=y,nxt[ed]=g[x],g[x]=ed;}
14 void init(){F(i,1,n)g[i]=0,vis[i]=0;ed=0,ret=0;}
15 
16 void dfs_size(int u,int fa)
17 {
18     size[u]=1,mx[u]=0;
19     for(int i=g[u];i;i=nxt[i])
20         if(v[i]!=fa&&!vis[v[i]])
21         {
22             dfs_size(v[i],u),size[u]+=size[v[i]];
23             if(size[v[i]]>mx[u])mx[u]=size[v[i]];
24         }
25 }
26 
27 void dfs_root(int r,int u,int fa)
28 {
29     if(size[r]-size[u]>mx[u])mx[u]=size[r]-size[u];
30     if(mx[u]<mi)mi=mx[u],root=u;
31     for(int i=g[u];i;i=nxt[i])
32         if(v[i]!=fa&&!vis[v[i]])
33             dfs_root(r,v[i],u);
34 }
35 
36 void dfs_dis(int u,int mi,int mx,int fa)
37 {
38     mi=min(mi,w[u]),mx=max(mx,w[u]);
39     if(mx<=mi+k)dis[++tot]=P(mi,mx);
40     for(int i=g[u];i;i=nxt[i])
41         if(v[i]!=fa&&!vis[v[i]])
42             dfs_dis(v[i],mi,mx,u);
43 }
44 
45 ll calc(int u,int mi,int mx)
46 {
47     ll ans=0;
48     tot=0,dfs_dis(u,mi,mx,0);
49     sort(dis+1,dis+tot+1);
50     F(i,1,tot)
51     {
52         int p=lower_bound(dis+1,dis+i+1,P(dis[i].second-k,0))-dis;
53         ans+=i-p;
54     }
55     return ans;
56 }
57 
58 void dfs(int u=1)
59 {
60     mi=n,dfs_size(u,0);
61     dfs_root(u,u,0);
62     ret+=calc(root,w[root],w[root]),vis[root]=1;
63     for(int i=g[root];i;i=nxt[i])
64         if(!vis[v[i]])ret-=calc(v[i],w[root],w[root]);
65     for(int i=g[root];i;i=nxt[i])
66         if(!vis[v[i]])dfs(v[i]);
67 }
68 
69 int main()
70 {
71     scanf("%d",&t);
72     while(t--)
73     {
74         scanf("%d%d",&n,&k);
75         init();
76         F(i,1,n)scanf("%d",w+i);
77         F(i,1,n-1)
78         {
79             int x,y;
80             scanf("%d%d",&x,&y);
81             adg(x,y),adg(y,x);
82         }
83         dfs(),printf("%lld\n",ret*2);
84     }
85     return 0;
86 }