[HDOJ5877]Weak Pair（DFS，线段树，离散化）

题目链接：http://acm.hdu.edu.cn/showproblem.php?pid=5877

题意：给一棵树和各点的权值a，求点对(u,v)个数，满足：1.u是v的祖先，2.a(u)*a(v)<=k。

对于这棵树，我们先存好树的结构。再离散化，最后dfs的时候往线段树里插点，那对应idx的值就是1。然后二分找不大于k/a[v]的下标，线段树统计计数就行了。换儿子的时候记得抹去上一个兄弟。

#include <bits/stdc++.h>
using namespace std;

#define lrt rt << 1
#define rrt rt << 1 | 1
typedef long long LL;
const int maxn = 200100;
int n, rt;
int in[maxn];
LL k, ret, a[maxn];
vector<int> G[maxn];
LL h[maxn];
int hcnt;
LL sum[maxn<<3];

inline int getid(LL x) {
  return lower_bound(h, h+hcnt, x) - h + 1;
}

void pushUP(int rt) {
  sum[rt] = sum[lrt] + sum[rrt];
}

void build(int l, int r, int rt) {
  sum[rt] = 0;
  if(l == r) return;
  int mid = (l + r) >> 1;
  build(l, mid, lrt);
  build(mid+1, r, rrt);
  pushUP(rt);
}

void update(int l, int r, int rt, int pos, LL val) {
  if(l == r) {
    sum[rt] += val;
    return;
  }
  int mid = (l + r) >> 1;
  if(pos <= mid) update(l, mid, lrt, pos, val);
  else update(mid+1, r, rrt, pos, val);
  pushUP(rt);
}

LL query(int L, int R, int l, int r, int rt) {
  if(l >= L && R >= r) return sum[rt];
  int mid = (l + r) >> 1;
  LL ret = 0;
  if(L <= mid) ret += query(L, R, l, mid, lrt);
  if(mid < R) ret += query(L, R, mid+1, r, rrt);
  return ret;
}

void dfs(int u) {
  int uu = getid(a[u]);
  int vv = getid(k/a[u]);
  ret += query(1, vv, 1, hcnt, 1);
  update(1, hcnt, 1, uu, 1);
  for(int i = 0; i < G[u].size(); i++) dfs(G[u][i]);
  update(1, hcnt, 1, uu, -1);
}

int main() {
  //freopen("in", "r", stdin);
  int T, u, v;
  scanf("%d", &T);
  while(T--) {
    scanf("%d %I64d",&n,&k);
    memset(in, 0, sizeof(in)); hcnt = 0; ret = 0;
    for(int i = 1; i <= n; i++) {
      scanf("%I64d", &a[i]);
      G[i].clear();
      h[hcnt++] = a[i]; h[hcnt++] = k / a[i];
    }
    sort(h, h+hcnt); hcnt = unique(h, h+hcnt) - h;
    build(1, hcnt, 1);
    for(int i = 0; i < n-1; i++) {
      scanf("%d %d",&u,&v);
      G[u].push_back(v);
      in[v]++;
    }
    for(int i = 1; i <= n; i++) if(!in[i]) dfs(i);
    printf("%I64d
", ret);
  }
  return 0;
}