HDU - 5877 Weak Pair （dfs+线段树）

Weak Pair

Time Limit: 4000/2000 MS (Java/Others)    Memory Limit: 262144/262144 K (Java/Others)
Total Submission(s): 5233    Accepted Submission(s): 1521
Problem Description

You are given a rooted tree of N nodes, labeled from 1 to N. To the ith node a non-negative value ai is assigned.An ordered pair of nodes (u,v) is said to be weak if
  (1) u is an ancestor of v (Note: In this problem a node u is not considered an ancestor of itself);
  (2) au×av≤k.

Can you find the number of weak pairs in the tree?

 Input

There are multiple cases in the data set.
  The first line of input contains an integer T denoting number of test cases.
  For each case, the first line contains two space-separated integers, N and k, respectively.
  The second line contains N space-separated integers, denoting a1 to aN.
  Each of the subsequent lines contains two space-separated integers defining an edge connecting nodes u and v , where node u is the parent of node v.

  Constrains: 
  
  1≤N≤105 
  
  0≤ai≤109 
  
  0≤k≤1018

 Output

For each test case, print a single integer on a single line denoting the number of weak pairs in the tree.

 Sample Input

1 2 3 1 2 1 2

 Sample Output

1

 Source

2016 ACM/ICPC Asia Regional Dalian Online

先离散化，在dfs，每到一个点，答案加上1到k/a[u]的区间和，然后a[u]对应的位置加1，对子树遍历结束后相应位置在减1，注意除0的情况

#include <bits/stdc++.h>
#define lson num << 1
#define rson num << 1 | 1
using namespace std;
typedef long long ll;
const int MAXN = 400015;
const ll INF = 1e19;
bool flag[MAXN];
vector<int> G[MAXN];
ll x[MAXN],X[MAXN];
int cnt;
ll ans,k;
struct node
{
    int l,r,sum;
}tree[MAXN << 2];
void build(int num,int l,int r)
{
    tree[num].l = l;
    tree[num].r = r;
    tree[num].sum = 0;
    if(l == r) return;
    int mid = (l + r) >> 1;
    build(lson,l,mid);
    build(rson,mid + 1,r);
}
int query(int num,int l,int r)
{
    if(tree[num].l == l && tree[num].r == r) {
        return tree[num].sum;
    }
    int sum = 0;
    int mid = (tree[num].l + tree[num].r) >> 1;
    if(r <= mid) sum += query(lson,l,r);
    else if(l > mid) sum += query(rson,l,r);
    else {
        sum += query(lson,l,mid);
        sum += query(rson,mid + 1,r);
    }
    return sum;
}
void update(int num,int pos,int val)
{
    if(tree[num].l == tree[num].r) {
        tree[num].sum += val;
        return;
    }
    int mid = (tree[num].l + tree[num].r) >> 1;
    if(pos <= mid) update(lson,pos,val);
    else update(rson,pos,val);
    tree[num].sum = tree[lson].sum + tree[rson].sum;
}
void DFS(int u)
{
    int r = lower_bound(X + 1,X + 1 + cnt,x[u] ?  k / x[u] : INF) - X;
    int pos = lower_bound(X + 1,X + 1 + cnt,x[u]) - X;
    ans += (query(1,1,r));
    update(1,pos,1);
    for(int i = 0; i < G[u].size(); i++) {
        DFS(G[u][i]);
    }
    update(1,pos,-1);
}
int main(void)
{
    int T,u,v,n,root;
    scanf("%d",&T);
    while(T--) {
        scanf("%d %lld",&n,&k);
        for(int i = 1; i <= n; i++) {
            scanf("%lld",&x[i]);
            X[i] = x[i];
            flag[i] = false;
            G[i].clear();
        }
        for(int i = (n + 1); i <= 2 * n; i++) {
            if(x[i - n]) X[i] = k / x[i - n];
            else X[i] = INF;
        }
        for(int i = 1; i < n; i++) {
            scanf("%d %d",&u,&v);
            G[u].push_back(v);
            flag[v] = 1;
        }
        for(int i = 1; i <= n; i++) {
            if(!flag[i]) {
                root = i;
                break;
            }
        }
        sort(X + 1,X + 1 + 2 * n);
        cnt = unique(X + 1,X + 1 + 2 * n) - (X + 1);
        build(1,1,cnt);
        ans = 0;
        DFS(root);
        printf("%lld\n",ans);
    }
    return 0;
}