【树分治】 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;}// headstruct 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;}