hdu 5877 Weak Pair dfs + 线段树（or树状数组）

//  hdu 5877 Weak Pair dfs + 线段树（or树状数组）
//
//  题目链接：http://acm.split.hdu.edu.cn/showproblem.php?pid=5877
//
//  题目大意：
//    n个节点到有根树，u是v的祖先，每个节点有个权值a[u].求
//    a[u] * a[v] <= k的有序对的数量。
//
//  解答：
//    我们可以进行dfs，这样可以将当前节点和它之前到祖先关联
//    起来，即，在它之前到一定是它到祖先。这样对于当前节点u
//    我们查找小于等于k/a[u]到u到祖先的个数x，将x累加即可得
//    到答案，这里我们可以用线段树或者树状数组。将a[u]的值
//    映射到线段树中，由于a[u]比较大，离散化一下，即可。
//
//  感悟：
//    这题，仔细想来，还是可以做的，只是真的好久没写代码了，
//    都生疏了，心中有着一丝到恐惧，现在慢慢捡起来，有段日子
//    觉得代码很恶心（我错了），但每次捡起来，都觉得很有趣，
//    这次，我将不打算放弃，将代码进行到底。内心成长。加油

#include <cstdio>
#include <iostream>
#include <cstring>
#include <stack>
#include <queue>
#include <algorithm>
using namespace std;
typedef long long ll;

const int MAXN = 1e5 + 8;
int a[MAXN];
int b[MAXN];

int in[MAXN];

int cnt;
ll sum = 0;
int n;

ll k;
vector<int> g[MAXN];

struct SegmentTree{

#define lson(x) (x << 1)
#define rson(x) (x << 1 | 1)

int sum[MAXN << 2];
int ql,qr;

void setq(int l,int r){
ql = l;
qr = r;
}

void setp(int l){
ql = l;
}

void init(){
memset(sum,0,sizeof(sum));
}

void push_up(int rt){
sum[rt] = sum[lson(rt)] + sum[rson(rt)];
}

void modify(int rt,int v,int L,int R){
if (L == R){
sum[rt] += v;
return ;
}

int M = (L + R) >> 1;
if (ql <= M)
modify(lson(rt),v,L,M);
else
modify(rson(rt),v,M+1,R);
push_up(rt);
}

int query(int rt,int L,int R){
if (ql <= L && R <= qr){
return sum[rt];
}
int M = (L + R) >> 1;
int res = 0;
if (ql <= M)
res += query(lson(rt),L,M);
if (qr > M)
res += query(rson(rt),M+1,R);
return res;
}

}it;

void dfs(int u,int fa){

it.setq(1,upper_bound(b + 1,b + cnt + 1,k / a[u]) - b - 1); // 查找区间[1.k/a[u])

// cout << k / a[u] << "----" << it.qr <<
//     "----" << it.query(1,1,cnt) << endl;
sum += it.query(1,1,cnt);
it.setp(lower_bound(b + 1,b + cnt + 1,a[u])- b); // 找到a[u]的位置
// cout << "before" << pos << " " << u << " " << a[u] << endl;
// cout << "sum: " << sum << endl;
if (!(int)g[u].size())
return;

it.modify(1,1,1,cnt); // 插入

for (int i = 0;i < g[u].size();i ++){
if (g[u][i] == fa)
continue;
dfs(g[u][i],u);
}

it.setp(lower_bound(b + 1,b + cnt + 1,a[u])- b);
it.modify(1,-1,1,cnt); // 删除（因为递归回来的时候，这个节点贡献值完
// 毕，即已经遍历完以它为父节点到子树）

}

int main(){
int t;
scanf("%d",&t);
while(t--){

scanf("%d%I64d\n",&n,&k);

for (int i = 1;i <= n;i ++) // 初始化
g[i].clear();
sum = 0;
it.init();

for (int i = 1;i <= n;i ++){

scanf("%d", a + i);
b[i] = a[i];
}

sort(b + 1,b + n + 1); // 排序，离散化的准备

cnt = unique(b + 1,b + n + 1) - b - 1; // 离散化

memset(in,0,sizeof(in));
for (int i = 1;i < n;i ++){
int u,v;
scanf("%d%d",&u,&v);
g[u].push_back(v);
in[v]++; // 入度为0到点为根
}
int rt = 0;

for (int i = 1;i <= n;i ++){
if (!in[i]){
rt = i;
break;
}
}

dfs(rt,-1);
printf("%I64d\n",sum);

}
return 0;
}