【树分治】 HDOJ 5314 Happy King

考虑经过根的路径，路径排序后用set统计答案。。。

#include <iostream>
#include <queue> 
#include <stack> 
#include <map> 
#include <set> 
#include <bitset> 
#include <cstdio> 
#include <algorithm> 
#include <cstring> 
#include <climits>
#include <cstdlib>
#include <cmath>
#include <time.h>
#define maxn 200005
#define maxm 400005
#define eps 1e-7
#define mod 1000000007
#define INF 0x3f3f3f3f
#define PI (acos(-1.0))
#define lowbit(x) (x&(-x))
#define mp make_pair
#define ls o<<1
#define rs o<<1 | 1
#define lson o<<1, L, mid 
#define rson o<<1 | 1, mid+1, R
#define pii pair<int, int>
#pragma comment(linker, "/STACK:16777216")
typedef long long LL;
typedef unsigned long long ULL;
//typedef int LL;
using namespace std;
LL qpow(LL a, LL b){LL res=1,base=a;while(b){if(b%2)res=res*base;base=base*base;b/=2;}return res;}
LL powmod(LL a, LL b){LL res=1,base=a;while(b){if(b%2)res=res*base%mod;base=base*base%mod;b/=2;}return res;}
// head

struct Edge
{
	int v;
	Edge *next;
}*H[maxn], E[maxm], *edges;

vector<pii> dis;
multiset<int> s;
multiset<int>::iterator it;
bool done[maxn];
int size[maxn];
int mx[maxn];
int a[maxn];
int n, T, root, nsize;
LL ans;

void addedges(int u, int v)
{
	edges->v = v;
	edges->next = H[u];
	H[u] = edges++;
}

void init()
{
	ans = 0;
	edges = E;
	memset(H, 0, sizeof H);
	memset(done, 0, sizeof done);
}

void getroot(int u, int fa)
{
	mx[u] = 0, size[u] = 1;
	for(Edge *e = H[u]; e; e = e->next) if(!done[e->v] && e->v != fa) {
		int v = e->v;
		getroot(v, u);
		size[u] += size[v];
		mx[u] = max(mx[u], size[v]);
	}
	mx[u] = max(mx[u], nsize - size[u]);
	if(mx[u] < mx[root]) root = u;
}

void getdis(int u, int fa, int maxv, int minv)
{
	dis.push_back(mp(maxv, minv));
	for(Edge *e = H[u]; e; e = e->next) if(!done[e->v] && e->v != fa) {
		int v = e->v;
		getdis(v, u, max(maxv, a[v]), min(minv, a[v]));
	}
}

void calc(int u, int flag, int maxv, int minv)
{
	int res = 0;
	dis.clear();
	s.clear();
	getdis(u, u, max(maxv, a[u]), min(minv, a[u]));
	
	sort(dis.begin(), dis.end());

	for(int i = 0; i < dis.size(); i++) {
		while(!s.empty() && dis[i].first - *(s.begin()) > T) {
			it = s.begin();
			s.erase(it);
		}
		if(!s.empty() && dis[i].first - dis[i].second <= T) ans += ((int)s.size()) * flag;
		s.insert(dis[i].second);
	}
}

void solve(int u)
{
	calc(u, 1, a[u], a[u]);
	done[u] = true;
	for(Edge *e = H[u]; e; e = e->next) if(!done[e->v]) {
		int v = e->v;
		calc(v, -1, a[u], a[u]);
		mx[0] = nsize = size[v];
		getroot(v, root = 0);
		solve(root);
	}
}

void work()
{
	int u, v;
	scanf("%d%d", &n, &T);
	for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
	
	for(int i = 1; i < n; i++) {
		scanf("%d%d", &u, &v);
		addedges(u, v);
		addedges(v, u);
	}
	
	ans = 0;
	mx[0] = nsize = n;
	getroot(1, root = 0);
	solve(root);
	printf("%I64d\n", (LL)2 * ans);
}

int main()
{
//	freopen("input", "r", stdin);
//	freopen("output", "w", stdout);

	int _;
	scanf("%d", &_);
	while(_--) {
		init();
		work();
	}
	
	return 0;
}